xView2 比賽冠軍程式碼解讀

钢之炼丹术师發表於2024-05-19
View

CSDN搬家失敗，手動匯出markdown後再匯入部落格園

程式碼地址：https://github.com/vdurnov/xview2_1st_place_solution

模型訓練中用到了混合精度訓練工具 Nvidia apex 和影像增強工具 imgaug

目錄

1、readme

權重檔案

資料清洗

資料處理

模型細節

2、程式碼結構

3、定位模型

3.1 資料集

3.2 模型結構

3.3 最佳化器

3.4 損失函式

4 分類模型

4.1 資料集

4.2 模型結構

4.3 損失函式

5 預測階段（待完善）

5.1 定位模型

5.2 分類模型

5.3 模型結果整合

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 模型結構

# 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.3 最佳化器

這裡採用了修正後的 Adam 最佳化器

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

大體意思是說現在常用的深度學習框架，在實現權重衰減的時候，都是使用直接在損失函式上 L2 正則來近似的。而實際上只有在使用 SGD 最佳化器的時候，L2 正則和權重衰減才等價。

在使用 Adam 最佳化器時，由於 L2 正則項裡面要除以 w 的平方，導致越大的權重實際衰減的越小，不符合預期，因此採用直接做權重衰減的方式來修正。

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 = {}

3.4 損失函式

損失函式這裡先介紹下作者自己定義的一個組合損失函式類

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)

裡面定義了多個損失函式，使用的時候給不同的 loss 賦予權重即可，如

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

對於定位模型來說，輸出就是一個單通道的影像，直接與 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

  2. 讀取定位輸出的 part1，做平均，歸一化

  3. msk_dmg 為每個通道上輸出最大值的位置，即把每個通道上的機率值轉化為 1~4 的標籤

  4. msk_loc 定位閾值篩選條件： _thr=[0.38,0.13,0.14]

    1. 定位大於閾值 0
    2. 或定位結果大於閾值 1 且 msk_dmg 對應位置大於 1 且小於 4（即 2~3 級破壞）
    3. 或定位結果大於閾值 2 且 msk_dmg 對應位置大於 1(即所有等級破壞）

  5. msk_dmg 更新為 msk_dmg 和 msk_loc 取交集

  6. 對於 2 級破壞，命名為_msk，如果有 2 級破壞 if _msk.sum()> 0，則對其進行膨脹，將膨脹後的_msk 與原始 msk_dmg==1 的位置做交集，即 2 級破壞膨脹後和 1 級破壞做交集，相交（重合）的地方設為 2 級破壞

相關文章