其他分享
首页 > 其他分享> > xView2比赛冠军代码解读

xView2比赛冠军代码解读

作者：互联网

代码地址：https://github.com/vdurnov/xview2_1st_place_solution

模型训练中用到了混合精度训练工具Nvidia apex和图像增强工具imgaug

目录

目录

1、readme

权重文件

数据清洗

数据处理

模型细节

2、代码结构

3、定位模型

3.1 数据增强

3.2 模型结构

3.3 优化器

3.4 损失函数

1、readme

权重文件

https://vdurnov.s3.amazonaws.com/xview2_1st_weights.zip

数据清洗

本次比赛的数据集非常完善，未发现有任何问题。使用json文件创建mask图像，将“un-classified”标签归到“no-damage”类别（create_masks.py）。

灾前和灾后图像，由于拍摄时的天底（nadir）不同，导致有微小的偏移。这个问题在模型层面被解决：

  1. 定位模型只使用灾前图像，以忽略灾后图像产生的这些偏移噪声。这里使用了Unet-like这种编码-解码结构的神经网络。
  2. 训练完成的定位模型转化为用于分类的孪生神经网络。如此，灾前和灾后的图像共享定位模型的权重，把最后一层解码层输出的特征concat来预测每个像素的损坏等级。这种方式使得神经网络使用相同的方式分别观察灾前和灾后图像，有助于忽略这些漂移。
  3. 使用5*5的kernel对分类mask做形态学膨胀，这使得预测结果更“bold”（理解为边界更宽一点），这有助于提升边界精度，并消除漂移和天底影响。

数据处理

 模型输入的尺寸从（448,448）到（736,736），越重的编码器结构采用越小的尺寸。训练所用的数据增强手段有：

推理阶段使用全尺寸图像（1024,1024），使用4中简单的测试阶段增强（原图，左右反转，上下翻转，180度旋转）

模型细节

所有模型训练集/验证集比例为9:1，每个模型使用3个随机数种子训练3次，保存最高验证集精度的checkpoint，结果保存在3个文件夹中。

定位模型使用了torchvision中预训练的4个encoder模型：

定位模型在灾前图像上训练，灾后图像只在极少数情况下作为额外的数据增强手段添加。

定位模型训练阶段参数：

分类模型使用对应的定位模型（和随机数种子）初始化。分类模型实际上是使用了整个定位模型的孪生神经网络，同时输入灾前和灾后图像。解码器最后一层的特征结合后用于分类。预训练的权重并不冻结。使用定位模型的预训练权重使得分类模型训练更快，精度更高。从灾前和灾后图像提取的特征在解码器的最后连接（bottleneck），这有助于防止过拟合，生成适用性更强的模型。

分类模型训练阶段参数：

所有checkpoint最后在整个训练集上做少次微调，使用低学习率和少量数据增强。

最终预测结果是把定位和分类模型的输出分别做平均。

定位模型对受损和未受损的类别使用不同的阈值（受损的更高）。

Pytorch预训练模型： https://github.com/Cadene/pretrained-models.pytorch

2、代码结构

首先看train.sh脚本。

echo "Creating masks..."
python create_masks.py
echo "Masks created"

echo "training seresnext50 localization model with seeds 0-2"
python train50_loc.py 0
python train50_loc.py 1
python train50_loc.py 2
python tune50_loc.py 0
python tune50_loc.py 1
python tune50_loc.py 2

该脚本使用3个不同的随机数种子将每个模型训练3次，然后微调，先训练定位模型，在训练分类模型，所有权重保存在weights文件夹中。

然后使用predict.sh脚本预测

python predict34_loc.py
python predict50_loc.py
python predict92_loc.py
python predict154_loc.py
python predict34cls.py 0
python predict34cls.py 1
python predict34cls.py 2
python predict50cls.py 0
python predict50cls.py 1
python predict50cls.py 2
python predict92cls.py 0
python predict92cls.py 1
python predict92cls.py 2
python predict154cls.py 0
python predict154cls.py 1
python predict154cls.py 2
echo "submission start!"
python create_submission.py
echo "submission created!"

使用所有模型进行定位和预测后，使用平均方式整合结果。

3、定位模型

下面以ResNet34_Unet为例讲解（train34_loc.py），其他模型除了结构以外几乎一样。首先看主函数，主要重点在于数据集、模型结构、优化器、损失函数

cv2.setNumThreads(0)
cv2.ocl.setUseOpenCL(False)

train_dirs = ['train', 'tier3']

models_folder = 'weights'

input_shape = (736, 736)

if __name__ == '__main__':
    t0 = timeit.default_timer()

    makedirs(models_folder, exist_ok=True)
    # 命令行接受一个参数，作为随机数种子
    seed = int(sys.argv[1])
    # vis_dev = sys.argv[2]

    # os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
    # os.environ["CUDA_VISIBLE_DEVICES"] = vis_dev

    cudnn.benchmark = True

    batch_size = 16
    val_batch_size = 8

    snapshot_name = 'res34_loc_{}_1'.format(seed)

    train_idxs, val_idxs = train_test_split(np.arange(len(all_files)), test_size=0.1, random_state=seed)

    np.random.seed(seed + 545)
    random.seed(seed + 454)

    steps_per_epoch = len(train_idxs) // batch_size
    validation_steps = len(val_idxs) // val_batch_size

    print('steps_per_epoch', steps_per_epoch, 'validation_steps', validation_steps)
    # 重点1：数据增强
    data_train = TrainData(train_idxs)
    val_train = ValData(val_idxs)

    train_data_loader = DataLoader(data_train, batch_size=batch_size, num_workers=6, shuffle=True, pin_memory=False, drop_last=True)
    val_data_loader = DataLoader(val_train, batch_size=val_batch_size, num_workers=6, shuffle=False, pin_memory=False)
    # 重点2：模型结构
    model = Res34_Unet_Loc()
    
    params = model.parameters()
    # 重点3：优化器
    optimizer = AdamW(params, lr=0.00015, weight_decay=1e-6)

    scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=[5, 11, 17, 25, 33, 47, 50, 60, 70, 90, 110, 130, 150, 170, 180, 190], gamma=0.5)

    model = nn.DataParallel(model).cuda()
    # 重点4：损失函数
    seg_loss = ComboLoss({'dice': 1.0, 'focal': 10.0}, per_image=False).cuda() #True

    best_score = 0
    _cnt = -1
    torch.cuda.empty_cache()
    for epoch in range(55):
        train_epoch(epoch, seg_loss, model, optimizer, scheduler, train_data_loader)
        if epoch % 2 == 0:
            _cnt += 1
            torch.cuda.empty_cache()
            best_score = evaluate_val(val_data_loader, best_score, model, snapshot_name, epoch)

    elapsed = timeit.default_timer() - t0
    print('Time: {:.3f} min'.format(elapsed / 60))

3.1 数据集

以训练集为例，添加注释

class TrainData(Dataset):
    def __init__(self, train_idxs):
        super().__init__()
        self.train_idxs = train_idxs
        self.elastic = iaa.ElasticTransformation(alpha=(0.25, 1.2), sigma=0.2)

    def __len__(self):
        return len(self.train_idxs)

    def __getitem__(self, idx):
        _idx = self.train_idxs[idx]

        fn = all_files[_idx]

        img = cv2.imread(fn, cv2.IMREAD_COLOR)
        # 少量使用灾后图像
        if random.random() > 0.985:
            img = cv2.imread(fn.replace('_pre_disaster', '_post_disaster'), cv2.IMREAD_COLOR)

        msk0 = cv2.imread(fn.replace('/images/', '/masks/'), cv2.IMREAD_UNCHANGED)
        # 水平翻转
        if random.random() > 0.5:
            img = img[::-1, ...]
            msk0 = msk0[::-1, ...]
        # 旋转
        if random.random() > 0.05:
            rot = random.randrange(4)
            if rot > 0:
                img = np.rot90(img, k=rot)
                msk0 = np.rot90(msk0, k=rot)
        # 仿射变换偏移
        if random.random() > 0.8:
            shift_pnt = (random.randint(-320, 320), random.randint(-320, 320))
            img = shift_image(img, shift_pnt)
            msk0 = shift_image(msk0, shift_pnt)
        # 旋转缩放
        if random.random() > 0.2:
            rot_pnt =  (img.shape[0] // 2 + random.randint(-320, 320), img.shape[1] // 2 + random.randint(-320, 320))
            scale = 0.9 + random.random() * 0.2
            angle = random.randint(0, 20) - 10
            if (angle != 0) or (scale != 1):
                img = rotate_image(img, angle, scale, rot_pnt)
                msk0 = rotate_image(msk0, angle, scale, rot_pnt)
        # 裁剪
        crop_size = input_shape[0]
        if random.random() > 0.3:
            crop_size = random.randint(int(input_shape[0] / 1.2), int(input_shape[0] / 0.8))

        bst_x0 = random.randint(0, img.shape[1] - crop_size)
        bst_y0 = random.randint(0, img.shape[0] - crop_size)
        bst_sc = -1
        try_cnt = random.randint(1, 5)
        for i in range(try_cnt):
            x0 = random.randint(0, img.shape[1] - crop_size)
            y0 = random.randint(0, img.shape[0] - crop_size)
            _sc = msk0[y0:y0+crop_size, x0:x0+crop_size].sum()
            if _sc > bst_sc:
                bst_sc = _sc
                bst_x0 = x0
                bst_y0 = y0
        x0 = bst_x0
        y0 = bst_y0
        img = img[y0:y0+crop_size, x0:x0+crop_size, :]
        msk0 = msk0[y0:y0+crop_size, x0:x0+crop_size]

        
        if crop_size != input_shape[0]:
            img = cv2.resize(img, input_shape, interpolation=cv2.INTER_LINEAR)
            msk0 = cv2.resize(msk0, input_shape, interpolation=cv2.INTER_LINEAR)
        # RGB通道变换
        if random.random() > 0.97:
            img = shift_channels(img, random.randint(-5, 5), random.randint(-5, 5), random.randint(-5, 5))
        elif random.random() > 0.97:
            img = change_hsv(img, random.randint(-5, 5), random.randint(-5, 5), random.randint(-5, 5))
        # 直方图、高斯噪声、滤波
        if random.random() > 0.93:
            if random.random() > 0.97:
                img = clahe(img)
            elif random.random() > 0.97:
                img = gauss_noise(img)
            elif random.random() > 0.97:
                img = cv2.blur(img, (3, 3))
        # 饱和度、亮度、对比度
        elif random.random() > 0.93:
            if random.random() > 0.97:
                img = saturation(img, 0.9 + random.random() * 0.2)
            elif random.random() > 0.97:
                img = brightness(img, 0.9 + random.random() * 0.2)
            elif random.random() > 0.97:
                img = contrast(img, 0.9 + random.random() * 0.2)
        # 弹性变换        
        if random.random() > 0.97:
            el_det = self.elastic.to_deterministic()
            img = el_det.augment_image(img)

        msk = msk0[..., np.newaxis]
        # msk二值化
        msk = (msk > 127) * 1
        # 像素归一化到（-1,1）
        img = preprocess_inputs(img)

        img = torch.from_numpy(img.transpose((2, 0, 1))).float()
        msk = torch.from_numpy(msk.transpose((2, 0, 1))).long()

        sample = {'img': img, 'msk': msk, 'fn': fn}
        return sample

3.2 模型结构

Res34_Unet_Loc类比较简单，就是一个普通的5层下采样Unet，其中1-5层为ResNet34的5个block，6-10层为上采样层

3.3 优化器

这里采用了修正后的Adam优化器

AdamW原文：https://arxiv.org/abs/1711.05101

大体意思是说现在常用的深度学习框架，在实现权重衰减的时候，都是使用直接在损失函数上L2正则来近似的。而实际上只有在使用SGD优化器的时候，L2正则和权重衰减才等价。

在使用Adam优化器时，由于L2正则项里面要除以w的平方，导致越大的权重实际衰减的越小，不符合预期，因此采用直接做权重衰减的方式来修正。

# w = w - wd * lr * w
if group['weight_decay'] != 0:
    p.data.add_(-group['weight_decay'] * group['lr'], p.data)

# w = w - lr * w.grad
p.data.addcdiv_(-step_size, exp_avg, denom)

3.4 损失函数

损失函数这里先介绍下作者自己定义的一个组合损失函数类

class ComboLoss(nn.Module):
    def __init__(self, weights, per_image=False):
        super().__init__()
        self.weights = weights
        self.bce = StableBCELoss()
        self.dice = DiceLoss(per_image=False)
        self.jaccard = JaccardLoss(per_image=False)
        self.lovasz = LovaszLoss(per_image=per_image)
        self.lovasz_sigmoid = LovaszLossSigmoid(per_image=per_image)
        self.focal = FocalLoss2d()
        self.mapping = {'bce': self.bce,
                        'dice': self.dice,
                        'focal': self.focal,
                        'jaccard': self.jaccard,
                        'lovasz': self.lovasz,
                        'lovasz_sigmoid': self.lovasz_sigmoid}
        self.expect_sigmoid = {'dice', 'focal', 'jaccard', 'lovasz_sigmoid'}
        self.values = {}

里面定义了多个损失函数，使用的时候给不同的loss赋予权重即可，如

seg_loss = ComboLoss({'dice': 1.0, 'focal': 10.0}, per_image=False).cuda()

对于定位模型来说，输出就是一个单通道的图像，直接与mask做像素级别的分类损失即可，后面分类模型会复杂一些。

4 分类模型

分类模型与定位模型相比，其实整体上差别不大，只不过是复制了2个定位模型构造孪生神经网络，同时输入灾前和灾后的图像，下面具体讲一下（train34_cls.py）。

4.1 数据集

数据增强方法与前面基本一致，主要区别在于mask的构造，前面的定位模型是一个单类别的图像分割（有无建筑），这里变成了对建筑进行更高细粒度的4类别分割。

        数据构造：img灾前，img2灾后

  1. msk0灾前
  2. lbl_msk1灾后1~4都有，原始mask
  3. msk1~msk4：分别对应灾后的1~4
  4. msk：msk0~msk4组合，形成5通道mask
  5. 对msk的1~4通道做形态学膨胀，然后调整多边形：
    1. 对于msk1层，有任何2~4级标签的地方均设为0
    2. 对于msk3层，有2级的像素点设为0
    3. 对于msk4层，有2级和3级的像素点均设为0
    4. 对于msk0层，有1~4级的像素点均设为1（0层用于定位，有建筑就为1）
  6. lbl_msk：1~4标签
  7. img和img2组合成6通道图像，归一化为-1到1
  8. 训练集lbl_msk = msk.argmax(axis=2)
  9. 验证集lbl_msk = msk[..., 1:].argmax(axis=2)

4.2 模型结构

 这里采用孪生神经网络，Res34_Unet_Double，网络结构跟前面单个的Res34_Unet完全一致，只是在输出层把两个分支整合

    def forward(self, x):
        dec10_0 = self.forward1(x[:, :3, :, :])
        dec10_1 = self.forward1(x[:, 3:, :, :])
        dec10 = torch.cat([dec10_0, dec10_1], 1)
        return self.res(dec10)

最后经过1*1卷积self.res输出了一个5通道的特征图。

在训练时，使用上一步定位模型的权重作为预训练权重。

4.3 损失函数

损失函数的设计跟定位模型一样，但是使用方法不同。

        loss0 = seg_loss(out[:, 0, ...], msks[:, 0, ...])
        loss1 = seg_loss(out[:, 1, ...], msks[:, 1, ...])
        loss2 = seg_loss(out[:, 2, ...], msks[:, 2, ...])
        loss3 = seg_loss(out[:, 3, ...], msks[:, 3, ...])
        loss4 = seg_loss(out[:, 4, ...], msks[:, 4, ...])

        loss = 0.05 * loss0 + 0.2 * loss1 + 0.8 * loss2 + 0.7 * loss3 + 0.4 * loss4

输出的5个通道分别和mask的5个通道做损失，其中loss0为灾前图像上的定位损失，1-4为灾后图像上4个级别的分类损失。

5 预测阶段（待完善）

5.1 定位模型

  1. 每张图片采样4次，每张图片做x,y,xy翻转，组合成batchsize=4的图像输入（4，3，1024，1024）
  2. 模型输出后在经过一个sigmoid层，输出4个msk，将每个mask按翻转方式还原，然后取平均值作为最终输出
  3. 保存为part1图像，单通道

5.2 分类模型

  1. 灾前灾后图像叠加，采样4次，组合成（4，6，1024，1024）
  2. 输出经过sigmoid，输出4个mask，做平均
  3. 由于输出是5通道，将前三个通道保存一次part1，后三个通道保存一次part2

5.3 模型结果整合

  1. 读取分类输出的part1和part2，重新组合成5通道mask，所有12个模型做平均，除以255归一化为0到1
  1. 读取定位输出的part1，做平均，归一化
  2. msk_dmg为每个通道上输出最大值的位置，即把每个通道上的概率值转化为1~4的标签
  1. msk_loc定位阈值筛选条件： _thr=[0.38,0.13,0.14]
    1. 定位大于阈值0
    2. 或定位结果大于阈值1且msk_dmg对应位置大于1且小于4（即2~3级破坏）
    3. 或定位结果大于阈值2且msk_dmg对应位置大于1(即所有等级破坏）
  2. msk_dmg更新为msk_dmg和msk_loc取交集
  3. 对于2级破坏，命名为_msk，如果有2级破坏  if _msk.sum() > 0，则对其进行膨胀，将膨胀后的_msk与原始msk_dmg==1的位置做交集，即2级破坏膨胀后和1级破坏做交集，相交（重合）的地方设为2级破坏

标签：py,img,模型,self,random,解读,msk,冠军,xView2
来源： https://blog.csdn.net/qq_26751117/article/details/111444916