(U-Net) pytorch로 Mirroring Extrapolate를 적용하여 학습하기

친환경입냄새 2022. 5. 18. 13:52

마지막 포스팅 후 한달만에 포스팅을 올리네요...

그동안 요것저것 너무 바빠서 구현을 해놓고 정리를 못했습니다...

지난번 U-Net 관련 포스팅에서 Mirroring Extrapolate를 구현해봤었는데

오늘은 Mirroring Extrapolate를 적용하여 U-Net 모델을 학습 시키는 부분을 포스팅하겠습니다.

Mirroring Extrapolate

U-Net: Convolutional Networks for Biomedical Image Segmentation

지난번 Mirroring Extrapolate는 아래와 같이 구현했었습니다.

import torch
import numpy as np

# tensor shape (1, x, y)
def mirroring_Extrapolate(img):
    # mirroring 92 pixel

    x = img.shape[1]
    y = img.shape[2]

    np_img = np.array(img)

    np_img = np_img[0]

    if x < 388:
        pad_x_left = (572 - x) / 2
        pad_x_right = (572 - x) / 2
    else:
        pad_x_left = 92
        pad_x_right = 388 - (x % 388) + 92

    if y < 388:
        pad_y_up = (572 - y) / 2
        pad_y_down = (572 - y) / 2
    else:
        pad_y_up = 92
        pad_y_down = 388 - (y % 388) + 92

    np_img = np.pad(np_img, ((pad_x_left, pad_x_right), (pad_y_up, pad_y_down)), 'reflect')

    np_img = np_img[:, :, np.newaxis]

    return torch.from_numpy(np_img.transpose((2, 0, 1)))

해당 코드로 'Mirroring Extrapolated'를 진행하면 아래와 같이 이미지가 변경됩니다.

Input, Label 이미지 둘 다 Mirroring Extrapolate를 적용하여 Input 이미지 한장에 학습을 4번 진행하였습니다.

-> 아래 이미지에서 노란색 사각형 4부분

-> Label 이미지도 똑같이 4부분을 Crop하여 학습 진행

Mirroring Extrapolate가 적용된 Input, Label 이미지에서 4등분 하는 코드는 아래와 같습니다.

U-Net 모델 학습용 데이터셋은 Input, Label 이미지가 512x512 사이즈 고정이라

과감하게 하드코딩을... 했습니다..

Mirroring Extrapolate가 적용된 Input, Label 이미지에서 4등분

# input : 960 x 960 / output 572 x 572 4개
def divide_input(img):
    left_top = torchvision.transforms.functional.crop(img, 0, 0, 572, 572)
    right_top = torchvision.transforms.functional.crop(img, 388, 0, 572, 572)
    left_bottom = torchvision.transforms.functional.crop(img, 0, 388, 572, 572)
    right_bottom = torchvision.transforms.functional.crop(img, 388, 388, 572, 572)
    return left_top, right_top, left_bottom, right_bottom

# input : 960 x 960 / output 388 x 388 4개
def divide_label(img):
    left_top = torchvision.transforms.functional.crop(img, 92, 92, 388, 388)
    right_top = torchvision.transforms.functional.crop(img, 480, 92, 388, 388)
    left_bottom = torchvision.transforms.functional.crop(img, 92, 480, 388, 388)
    right_bottom = torchvision.transforms.functional.crop(img, 480, 480, 388, 388)
    return left_top, right_top, left_bottom, right_bottom

Mirroring Extrapolate는 Dataset을 불러오는 과정에서 적용을 시킨 후 아래와 같이 학습을 진행하였습니다.

batch_size = 4
lr = 1e-4
num_epoch = 100

dataset_train = Dataset(data_path + '/U-Net_train', train_transform, 'Mirroring')

loader_train = DataLoader(dataset_train, batch_size=batch_size, shuffle=True)

fn_loss = nn.BCEWithLogitsLoss().to(device)

optim = torch.optim.Adam(net.parameters(), lr=lr)

start_epoch = 0

for epoch in range(start_epoch + 1, num_epoch + 1):
    net.train()
    loss_arr = []
    print('epoch : ', epoch)
    for batch, data in enumerate(loader_train, 1):
        # forward
        label = data['label'].to(device) 
        inputs = data['image'].to(device)

        input_0, input_1, input_2, input_3 = divide_input(inputs)
        label_0, label_1, label_2, label_3 = divide_label(label)

        input_list = []
        label_list = []

        input_list.append(input_0)
        input_list.append(input_1)
        input_list.append(input_2)
        input_list.append(input_3)

        label_list.append(label_0)
        label_list.append(label_1)
        label_list.append(label_2)
        label_list.append(label_3)

        for i in range(4):
            print('input shape : ', input_list[i].shape)
            print('label shape : ', label_list[i].shape)
            output = net(input_list[i])

            print('output shape : ', output.shape)

            # backward
            optim.zero_grad()  # gradient 초기화

            loss = fn_loss(output, label_list[i])
            loss.backward()
            optim.step()

            # save loss
            loss_arr += [loss.item()]
            print('train : epoch %04d / %04d | Batch %04d \ %04d | divide %d \ 4 | Loss %.2f' % (
                epoch, num_epoch, batch, num_train_for_epoch, i, np.mean(loss_arr)))

위 방식대로 학습을 진행하면 1 Batch에 총 4번 학습을 진행하게 됩니다.

Epoch : 1, Batch : 1_1