代码地址:https://github.com/vdurnov/xview2_1st_place_solution
模型训练中用到了混合精度训练工具Nvidia apex和图像增强工具imgaug
目录
目录
1、readme
权重文件
https://vdurnov.s3.amazonaws.com/xview2_1st_weights.zip
数据清洗
本次比赛的数据集非常完善,未发现有任何问题。使用json文件创建mask图像,将“un-classified”标签归到“no-damage”类别(create_masks.py)。
灾前和灾后图像,由于拍摄时的天底(nadir)不同,导致有微小的偏移。这个问题在模型层面被解决:
- 定位模型只使用灾前图像,以忽略灾后图像产生的这些偏移噪声。这里使用了Unet-like这种编码-解码结构的神经网络。
- 训练完成的定位模型转化为用于分类的孪生神经网络。如此,灾前和灾后的图像共享定位模型的权重,把最后一层解码层输出的特征concat来预测每个像素的损坏等级。这种方式使得神经网络使用相同的方式分别观察灾前和灾后图像,有助于忽略这些漂移。
- 使用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灾后
- msk0灾前
- lbl_msk1灾后1~4都有,原始mask
- msk1~msk4:分别对应灾后的1~4
- msk:msk0~msk4组合,形成5通道mask
- 对msk的1~4通道做形态学膨胀,然后调整多边形:
- 对于msk1层,有任何2~4级标签的地方均设为0
- 对于msk3层,有2级的像素点设为0
- 对于msk4层,有2级和3级的像素点均设为0
- 对于msk0层,有1~4级的像素点均设为1(0层用于定位,有建筑就为1)
- lbl_msk:1~4标签
- img和img2组合成6通道图像,归一化为-1到1
- 训练集lbl_msk = msk.argmax(axis=2)
- 验证集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 定位模型
- 每张图片采样4次,每张图片做x,y,xy翻转,组合成batchsize=4的图像输入(4,3,1024,1024)
- 模型输出后在经过一个sigmoid层,输出4个msk,将每个mask按翻转方式还原,然后取平均值作为最终输出
- 保存为part1图像,单通道
5.2 分类模型
- 灾前灾后图像叠加,采样4次,组合成(4,6,1024,1024)
- 输出经过sigmoid,输出4个mask,做平均
- 由于输出是5通道,将前三个通道保存一次part1,后三个通道保存一次part2
5.3 模型结果整合
- 读取分类输出的part1和part2,重新组合成5通道mask,所有12个模型做平均,除以255归一化为0到1
- 读取定位输出的part1,做平均,归一化
- msk_dmg为每个通道上输出最大值的位置,即把每个通道上的概率值转化为1~4的标签
- msk_loc定位阈值筛选条件: _thr=[0.38,0.13,0.14]
- 定位大于阈值0
- 或定位结果大于阈值1且msk_dmg对应位置大于1且小于4(即2~3级破坏)
- 或定位结果大于阈值2且msk_dmg对应位置大于1(即所有等级破坏)
- msk_dmg更新为msk_dmg和msk_loc取交集
- 对于2级破坏,命名为_msk,如果有2级破坏 if _msk.sum() > 0,则对其进行膨胀,将膨胀后的_msk与原始msk_dmg==1的位置做交集,即2级破坏膨胀后和1级破坏做交集,相交(重合)的地方设为2级破坏