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),越重的编码器结构采用越小的尺寸。训练所用的数据增强手段有:

  • 翻转(often)
  • 旋转(often)
  • 缩放(often)
  • 颜色shift(rare)
  • 直方图、模糊、噪声(rare)
  • 饱和度、亮度、对比度(rare)
  • 弹性变换(rare)

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

模型细节

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

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

  • ResNet34
  • se_resnext50_32x4d
  • SeNet154
  • Dpn92

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

定位模型训练阶段参数:

  • 损失函数:Dice+Focal
  • 验证指标:Dice
  • 优化器:AdamW

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

分类模型训练阶段参数:

  • 损失函数:Dice+Focal+交叉熵。交叉熵损失函数中2-4级破坏的系数更大一些。
  • 验证指标:比赛提供的metric
  • 优化器:AdamW
  • 采样:2-4级破坏采样2次,对其更关注

所有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级破坏

猜你喜欢

转载自blog.csdn.net/qq_26751117/article/details/111444916