【U-Net】Semantic Segmentationの学習フェーズの実装 - 山下誠二郎のITとフォートナイトと野球ブログ

前書き

Kaggle（やってない）や研究の勉強ついでに。
ネットワークモデルは自分で構築するより、
外から取ってきて自分の実現したい内容に合わせるかが重要です。

概要

前回dataloaderの作成まで行いました。
今回はU-Netで

melheaven.hatenadiary.jp

今回はU-NetをPytorchで実装してみようと思います。
と言っても一から実装するわけでなく、以下の記事からコードを取得してきました。

github.com

U-Netの論文・モデル構造

以下はU-Netの論文(Arxiv)と参考記事。
U-NetとはFCN(Fully Convolution Network)の一つであり、
物体検出を画素毎に行ってくれるSemantic Segmentationを実現してくれます。

https://arxiv.org/pdf/1612.05360.pdf

U-Netのモデルの仕組みをわかりやすく

論文読んでも意味わかんない方はこちらがおすすめです。

www.acceluniverse.com

Script

モデル読み込み

まずgpuが使用可能か確認。今回はgpuを使用します。
ResNetUnetというベースがResNetのモデルを読み込みます。
ラベル数は対象オブジェクトとその他の2つです。

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = ResNetUNet(n_class=2)
model = model.to(device)

モデル確認

summaryでモデルの層をインラインで確認できます。
今回は入力画像に224×224の3channel画像とします。

!pip install torch-summary
# check keras-like model summary using torchsummary
from torchsummary import summary
summary(model, input_size=(3, 224, 224))

Trainingの流れ

Training：訓練を行う関数を自作します。
損失誤差を記録しながら学習を進めます。

def train_model(model, optimizer, scheduler, num_epochs=25):
    best_model_wts = copy.deepcopy(model.state_dict())
    best_loss = 1e10

   #****************
   # epoch毎に学習箇所(次で説明)
   #****************

    print('Best val loss: {:4f}'.format(best_loss))

    # load best model weights
    model.load_state_dict(best_model_wts)
    return model

epoch毎に学習箇所の中身

エポック数（1つの訓練データを何回繰り返し学習したか）と学習時間の表示。

    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
        print('-' * 10)

        since = time.time()

       #****************
       # 訓練・検証に分けて実行箇所(次で説明)
　　#****************

        time_elapsed = time.time() - since
        print('{:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))

訓練・検証に分けて実行箇所の中身

1epochあたり訓練・検証を繰り返しながら学習します。
前回の記事で訓練と検証に使用するデータは別々に分けて生成しました。
用意したモデルで訓練モード・検証モードに切り替えられます。

Schedulerは何をやってるかと言うと、
最適解を飛び越してしまわないように、
学習率をEpoch毎に調節しながら学習できる機能です。
あとで紹介しますが今回はStepLRを使用しています。

for phase in ['train', 'val']:
            if phase == 'train':
                scheduler.step()
                for param_group in optimizer.param_groups:
                    print("LR", param_group['lr'])

                model.train()  # Set model to training mode
            else:
                model.eval()   # Set model to evaluate mode

            metrics = defaultdict(float)
            epoch_samples = 0

            #****************
            # 学習の本筋部分(次で説明)
　　     #****************

　　　 print_metrics(metrics, epoch_samples, phase)
            epoch_loss = metrics['loss'] / epoch_samples

学習の本筋部分の中身

ここを読むと学習過程でデータの流れがざっと理解できると思います。

Tensor演算をcpu or gpuで設定
Optimizerの初期化
順伝播
損失関数の計算
誤差逆伝播（訓練時のみ）
Optimizerの更新

具体的にはこんな流れ

パラメータを用いてモデル(model)で計算を行い
出力値(output)と実際の値(labels)から誤差(loss)を計算
誤差関数(loss)のbackwardメソッドを呼び出し
optimizerの更新

for inputs, labels in dataloaders_dict[phase]:
                inputs = inputs.to(device)
                labels = labels.to(device)

                # zero the parameter gradients
                optimizer.zero_grad()

                # forward
                # track history if only in train
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    loss = calc_loss(outputs, labels, metrics)

                    # backward + optimize only if in training phase
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

Training関数の呼び出し

現在紹介していないのはOptimizerとScheduler、損失関数です。
これはTraining関数の呼び出し前段階で関数を作成するなり、オプションとして設定します。

model = train_model(model, optimizer_ft, exp_lr_scheduler, num_epochs=100)

感想

損失関数が呼び出すbackward()などの中身を見ていきたいと思います。
Optimizerや損失関数などを次回は紹介できたらと思います。