ChainerでSegNetとU-Net

Chainerを使ってSemantic Segmentationをやってみよう、ということで、SegNetを実装してみました（実はFCNもやってみたけど、なんかうまくいかなかった）。Semantic Segmentationというのは、ピクセル単位に識別をやるということで、SegNetは代表的なSemantic Segmentationのモデルの一つです。SegNetは2016年にVijay Badrinarayananらによって提案されました(原論文はこちら)。いまではSegNetより遥かに正確にSegmentationが可能なモデルもたくさん提唱されていますが、実装力を試してみるのによい課題だと思って挑戦してみました。SegNetについての詳しい説明は@cyberailabさんの記事にまとまっているのでここでは触れません。

またここでは、ガチのSegNetではなく、その廉価版のSegNet Basicを実装＆実行してみました。なお、実装に際しては、SegNetをKerasで実装した@uni-3さんの記事の記事をとても参考にしました。私の記事は@uni-3さんの記事の内容をchainerに移植したプラスアルファくらいの内容です。

以下の全コードはgitにあげています。

環境

python 3.5.2
chainer 5.1.0
Ubuntu 16.04
GeForce GTX 1050

データ

CamVidというデータセットを用いました。
https://github.com/alexgkendall/SegNet-Tutorial
リポジトリのCamVid内のデータを使います。クラス数は12です。

annotフォルダの中には真っ黒な画像が入っていますが、これをPILなどで開いてみると各セルに0から11までの数字が入っています。

ちなみに、このデータセットのみでは十分な汎化性能を出すことができません。もっと多くの訓練データが必要です。

2019/4/16 追記

U-Netの実装とsegmentationの結果を末尾に書き加えました。

実装

また、以下のコードを読む際にはpngのインデックスカラーについて知っておくとスムーズに理解が進むと思います。

下準備

まずは、諸設定を行うconfig.pyファイルを作ります。
dropoutをTrueにすると、学習に時間がかかりますが過学習を抑制できます。img_normalizationがTrueの場合は@uni-3さんの記事のdataset.pyで行われていた操作と同じことを前処理としてRGB画像に施します。LRNは局所的応答正規化のことで、TrueにするとネットワークにRGB画像を読み込む際に局所応答正規化を実行します。

GPUメモリに余裕がある場合はbatch_sizeを大きくするとよいと思います。

config.py

max_epoch = 101
batch_size = 1
n_class = 12
img_size = (360, 480)

train_img_path = "dataset/train"
train_gt_path = "dataset/trainannot"
valid_img_path = "dataset/val"
valid_gt_path = "dataset/valannot"
test_img_path = "dataset/test"
test_gt_path = "dataset/testannot"
save_predicted_path = "predicted_imgs"
save_model_path = "saved_models"
img_normalization = True
LRN = True
dropout = True

次に必要に応じて、dataset内の画像や画像名を参照するloader.pyを作ります。
Chainerでは、CNNに(バッチサイズ, チャンネル数, 縦, 横)の4次元配列を食わせるため、その準備として、(縦,横,チャンネル数)という並びの配列を(チャンネル数, 縦, 横)という順番にtransposeすることに注意してください。

loader.py

import os
import cv2
import numpy as np
import config

img_size = config.img_size
train_img_path = config.train_img_path
valid_img_path = config.valid_img_path
train_gt_path = config.train_gt_path
valid_gt_path = config.valid_gt_path

def img_loader(imgs_names, valid=False):
    num_imgs = len(imgs_names)

    imgs = np.empty((num_imgs, 3, img_size[0], img_size[1]), dtype="float32")
    for i in range(num_imgs):
        if not valid:
            img = cv2.imread(os.path.join(train_img_path, imgs_names[i]))
        else:
            img = cv2.imread(os.path.join(valid_img_path, imgs_names[i]))
        img = img.transpose(2,0,1)
        img = img.astype("float32")
        imgs[i] = img

    gts = np.empty((num_imgs, img_size[0], img_size[1]), dtype="int32")
    for i in range(num_imgs):
        if not valid:
            gt = cv2.imread(os.path.join(train_gt_path, imgs_names[i]))
        else:
            gt = cv2.imread(os.path.join(valid_gt_path, imgs_names[i]))
        gt = gt.astype("int32")
        gt = gt[:,:,0] #gt[:,:,0] == gt[:,:,1] == gt[:,:,2]
        gt = gt.reshape(img_size[0], img_size[1]) #which is better it or flatten()?
        gts[i] = gt

    return imgs, gts

def img_names_loader():
    imgs_names = os.listdir(train_img_path)
    imgs_names.sort()
    imgs_names = np.asarray(imgs_names)

    gt_names = os.listdir(train_gt_path)
    gt_names.sort()
    gt_names = np.asarray(gt_names)

    valid_imgs_names = os.listdir(valid_img_path)
    valid_imgs_names.sort()
    valid_imgs_names = np.asarray(valid_imgs_names)

    valid_gt_names = os.listdir(valid_gt_path)
    valid_gt_names.sort()
    valid_gt_names = np.asarray(valid_gt_names)

    return imgs_names, gt_names, valid_imgs_names, valid_gt_names

次にconfig.pyでimg_normalization = Trueとした場合に実行されるスクリプトを書きます。opencvでヒストグラム平坦化を各チャンネル毎に行います。

normalize.py

import cv2
import os
import config
import numpy as np

img_size = config.img_size

def normalize(imgs):
    num_imgs = len(imgs)
    imgs = imgs.astype("uint8")
    normalized_imgs = np.zeros((num_imgs, 3, img_size[0], img_size[1]), dtype="float32")
    for i in range(num_imgs):
        normalized_imgs[i,0,:,:] = cv2.equalizeHist(imgs[i,0,:,:])
        normalized_imgs[i,1,:,:] = cv2.equalizeHist(imgs[i,1,:,:])
        normalized_imgs[i,2,:,:] = cv2.equalizeHist(imgs[i,2,:,:])

    return normalized_imgs

モデル

SegNet Basicの本体をchainerで記述します。config.pyでLRN = Trueとすると、局所応答正規化を行うF.local_response_normalizationが入力rgb画像xに行われます。dropout = True とするとネットワークの中央でドロップアウトを生じさせます。これらは、原論文に忠実に加えました。

network.py

import chainer
import chainer.links as L 
import chainer.functions as F
import config
from chainer import Chain

n_class = config.n_class
LRN = config.LRN
dropout = config.dropout

class SegNet(Chain):

    def __init__(self, n_class=n_class):
        super().__init__()
        with self.init_scope():

            self.n_class = n_class

            self.enco1_1 = L.Convolution2D(None,  64, ksize=3, stride=1, pad=1) #image size will not change
            self.enco2_1 = L.Convolution2D(None, 128, ksize=3, stride=1, pad=1)
            self.enco3_1 = L.Convolution2D(None, 256, ksize=3, stride=1, pad=1)
            self.enco4_1 = L.Convolution2D(None, 512, ksize=3, stride=1, pad=1)

            self.deco4_1 = L.Convolution2D(None, 512, ksize=3, stride=1, pad=1)
            self.deco3_1 = L.Convolution2D(None, 256, ksize=3, stride=1, pad=1)
            self.deco2_1 = L.Convolution2D(None, 128, ksize=3, stride=1, pad=1)
            self.deco1_1 = L.Convolution2D(None,  64, ksize=3, stride=1, pad=1)
            self.deco0_1 = L.Convolution2D(None, n_class, ksize=1)

            self.bn_enco1_1 = L.BatchNormalization( 64)
            self.bn_enco2_1 = L.BatchNormalization(128)
            self.bn_enco3_1 = L.BatchNormalization(256)
            self.bn_enco4_1 = L.BatchNormalization(512)

            self.bn_deco4_1 = L.BatchNormalization(512)
            self.bn_deco3_1 = L.BatchNormalization(256)
            self.bn_deco2_1 = L.BatchNormalization(128)
            self.bn_deco1_1 = L.BatchNormalization( 64)

    def __call__(self, x): #x = (batchsize, 3, 360, 480)
        if LRN:
            x = F.local_response_normalization(x) #Needed for preventing from overfitting

        h = F.relu(self.bn_enco1_1(self.enco1_1(x)))
        h, idx1 = F.max_pooling_2d(h, 2, stride=2, return_indices=True)

        h = F.relu(self.bn_enco2_1(self.enco2_1(h)))
        h, idx2 = F.max_pooling_2d(h, 2, stride=2, return_indices=True) 

        h = F.relu(self.bn_enco3_1(self.enco3_1(h)))
        h, idx3 = F.max_pooling_2d(h, 2, stride=2, return_indices=True) 

        h = F.relu(self.bn_enco4_1(self.enco4_1(h)))
        if dropout:
            h = F.dropout(h)
        h, idx4 = F.max_pooling_2d(h, 2, stride=2, return_indices=True) 

        h = F.relu(self.bn_deco4_1(self.deco4_1(h)))
        h = F.unpooling_2d(h, ksize=2, stride=2, outsize=(idx3.shape[2], idx3.shape[3]))

        h = F.relu(self.bn_deco3_1(self.deco3_1(h)))
        h = F.unpooling_2d(h, ksize=2, stride=2, outsize=(idx2.shape[2], idx2.shape[3]))

        h = F.relu(self.bn_deco2_1(self.deco2_1(h)))
        h = F.unpooling_2d(h, ksize=2, stride=2, outsize=(idx1.shape[2], idx1.shape[3]))

        h = F.relu(self.bn_deco1_1(self.deco1_1(h)))
        h = F.unpooling_2d(h, ksize=2, stride=2, outsize=(x.shape[2], x.shape[3]))

        h = self.deco0_1(h)

        return h

訓練＆テスト

Trainを行うtrain.pyを作ります。loss関数は原論文にあるようにweight balancedされたsoftmax cross entropyを採用します。ここで唐突にclass_weighting配列が登場しますが、これは@uni-3さんの記事のtrain.pyから引用しました。この数字の根拠は分かりません。どのように求めたのか誰か知っていれば教えてください。

デフォルトで訓練中のモデルは10 epoch毎に.h5ファイルとしてsaved_modelに保存されます。また、同じく10 epoch毎にvalidation loss と training loss の変化を折れ線グラフで表したloss_curve.pngが保存されます。

私はtrainerなどを使って、コードを抽象化するのが苦手なのでベタ書きしています。

train.py

import os
import cv2
import time
import numpy as np
import cupy as xp
import matplotlib.pyplot as plt
from chainer import functions as F
from chainer import Variable, optimizers, serializers

import network
import config
import loader
import normalize

max_epoch = config.max_epoch
batch_size = config.batch_size
n_class = config.n_class
save_model_path = config.save_model_path
img_normalization = config.img_normalization

class_weighting = xp.array([0.2595, 0.1826, 4.5640, 0.1417, 0.5051, 0.3826, 9.6446, 1.8418, 6.6823, 6.2478, 3.0, 7.3614], dtype="float32")

def draw_loss_curve(t_loss, v_loss):
    t = np.arange(0, len(t_loss))
    plt.figure(figsize=(8,8))
    plt.plot(t, t_loss, label="train", color="b")
    plt.plot(t, v_loss, label="valid", color="r")
    plt.legend(fontsize=18)
    plt.xlabel("Epoch", fontsize=18)
    plt.ylabel("Loss", fontsize=18)
    plt.title("Loss Curve", fontsize=18)
    plt.ylim(0, 2.5)
    plt.savefig("loss_curve.png")
    plt.clf()

def train():    
    if not os.path.exists(save_model_path):
        os.makedirs(save_model_path)
        print("Made save folder")

    imgs_names, gt_names, valid_imgs_names, valid_gt_names = loader.img_names_loader()

    model = network.SegNet(n_class = n_class)
    model.to_gpu(0)
    optimizer = optimizers.Adam().setup(model)

    train_loss_recode = []
    valid_loss_recode = []

    N = len(imgs_names) 
    M = len(valid_imgs_names)
    perm = np.random.permutation(N)
    perm_valid = np.random.permutation(M)
    start_time = time.time()
    for epoch in range(max_epoch):
        losses = []
        for i in range(0, N, batch_size):
            imgs_names_batch = imgs_names[perm[i:i + batch_size]]
            imgs_batch, gt_batch = loader.img_loader(imgs_names_batch)

            if img_normalization:
                imgs_batch = normalize.normalize(imgs_batch)

            imgs_batch = Variable(imgs_batch)
            imgs_batch.to_gpu(0)
            gt_batch = Variable(gt_batch)
            gt_batch.to_gpu(0)

            model.cleargrads()
            t = model(imgs_batch)

            loss = F.softmax_cross_entropy(t, gt_batch, class_weight=class_weighting)
            loss.backward()
            optimizer.update()

            losses.append(loss.data)

        valid_losses = []
        for i in range(0, M, batch_size):
            valid_names_batch = valid_imgs_names[perm_valid[i:i + batch_size]]
            valid_imgs_batch, valid_gt_batch = loader.img_loader(valid_names_batch, valid=True)
            if img_normalization:
                valid_imgs_batch = normalize.normalize(valid_imgs_batch)

            valid_imgs_batch = Variable(valid_imgs_batch)
            valid_imgs_batch.to_gpu(0)
            valid_gt_batch = Variable(valid_gt_batch)
            valid_gt_batch.to_gpu(0)

            model.cleargrads()
            t = model(valid_imgs_batch)
            valid_loss = F.softmax_cross_entropy(t, valid_gt_batch, class_weight=class_weighting)
            valid_losses.append(valid_loss.data)

        train_loss_recode.append(sum(losses)/len(losses))
        valid_loss_recode.append(sum(valid_losses)/len(valid_losses))
        print("epoch:{0:}\t train loss:{1:.5f}\t valid loss:{2:.5f}\t time:{3:.2f}[sec]".format(epoch, float(train_loss_recode[-1]), float(valid_loss_recode[-1]), time.time()-start_time))

        if epoch % 10 == 0:
            draw_loss_curve(train_loss_recode, valid_loss_recode)
            serializers.save_hdf5(os.path.join(save_model_path, "SegNet_{}.h5".format(epoch)), model)
            print("saved model")

if __name__ == '__main__':
    train()

次にテスト用コードを書きます。dataset/train内にある画像を順番にトレーニング済みのモデルに突っ込んだ結果、得られた1チャンネル2次元画像を着色して保存します。そのあ後に、混合行列を作成します。overfittingしたモデルでテストすることを避けるために、デフォルトでは20epoch時に保存されたモデルを用いてテストを行います。

test.py

import cv2
import os
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
from chainer import serializers, Variable
from chainer.cuda import to_cpu 
from chainer import functions as F

import network
import config
import normalize
import confusion_matrix as cm

test_img_path = config.test_img_path
test_gt_path = config.test_gt_path
save_model_path = config.save_model_path
save_predicted_path = config.save_predicted_path
img_size = config.img_size
n_class = config.n_class
max_epoch = config.max_epoch
img_normalization = config.img_normalization

palette = [
    [128,128,128],
    [128,0,0],
    [192,192,128],
    [128,64,128],
    [255,69,0],
    [60,40,222],
    [128,128,0],
    [192,128,128],
    [64,64,128],
    [64,0,128],
    [64,64,0],
    [0,128,192],
    [0,0,0]]

def draw_png(img):
    seg_img = np.zeros((img_size[0], img_size[1], 3))
    for i in range(img_size[0]):
        for j in range(img_size[1]):
            color_number = img[i,j]
            seg_img[i,j,0] = palette[color_number][0]
            seg_img[i,j,1] = palette[color_number][1]
            seg_img[i,j,2] = palette[color_number][2]
    seg_img = Image.fromarray(np.uint8(seg_img))

    return seg_img
    #When you use cv2, RGB will be BGR. TAKE CARE IT.

def main():
    model = network.SegNet()
    model.to_gpu(0)
    serializers.load_hdf5(os.path.join(save_model_path, "SegNet_{}.h5".format(20)), model)

    if not os.path.exists(save_predicted_path):
        os.makedirs(save_predicted_path)
        print("Made save folder")

    test_names = os.listdir(test_img_path)
    test_names.sort()
    num_imgs = len(test_names)
    imgs_eval  = np.empty((num_imgs, img_size[0], img_size[1]))

    for i in range(num_imgs):
        img = cv2.imread(os.path.join(test_img_path, test_names[i]))
        img = img.transpose(2,0,1)
        img = img.reshape((1, 3, img_size[0], img_size[1]))
        img = img.astype("float32")
        if img_normalization:
            img = normalize.normalize(img)
        img = Variable(img)
        img.to_gpu(0)
        seg_img = model(img)
        seg_img.to_cpu()
        seg_img = seg_img.reshape((n_class, img_size[0], img_size[1]))
        seg_img = F.argmax(seg_img, axis=0).array
        colored_seg_img = draw_png(seg_img)
        colored_seg_img.save(os.path.join(save_predicted_path, test_names[i]))
        print("{} is saved".format(test_names[i]))
        imgs_eval[i] = seg_img

    confusion_matrix, acc = cm.make_confusion_matrix(imgs_eval)
    cm.show_confusion_matrix(confusion_matrix, acc)

if __name__ == '__main__':
    main()

評価

評価用に混合行列をつくります。この辺は好みですが、sklearnを使ってもよいと思います。私はフルスクラッチしました。

confusion_matrix.py

import cv2
import os
import network
import config
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
from chainer import serializers, Variable
from chainer.cuda import to_cpu 
from chainer import functions as F

test_gt_path = config.test_gt_path
img_size = config.img_size
n_class = config.n_class

def make_confusion_matrix(pre_imgs):
    pre_imgs = pre_imgs.astype("int32")
    gt_names = os.listdir(test_gt_path)
    gt_names.sort()

    num_gts = len(gt_names)
    gt_imgs = np.empty((num_gts, img_size[0], img_size[1]), dtype="int32")
    for i in range(num_gts):
        gt_img = cv2.imread(os.path.join(test_gt_path, gt_names[i]))
        gt_img = gt_img[:, :, 0]
        gt_imgs[i] = gt_img

    confusion_matrix = np.zeros((n_class, n_class), dtype="float32")
    print("making confusion matrix...")
    for i in range(len(pre_imgs)):
        for j in range(img_size[0]):
            for k in range(img_size[1]):
                confusion_matrix[gt_imgs[i,j,k], pre_imgs[i,j,k]] += 1

    acc = np.trace(confusion_matrix)/confusion_matrix.sum()
    acc = round(acc, 3)

    for i in range(n_class):
        if confusion_matrix[i].sum() != 0:
            confusion_matrix[i] = confusion_matrix[i] / confusion_matrix[i].sum()

    return confusion_matrix, acc

def show_confusion_matrix(confusion_matrix, acc):
    fig, ax = plt.subplots(figsize=(11, 8))
    heatmap = ax.pcolor(confusion_matrix, cmap=plt.cm.Blues)

    labels = ["Sky", "Building", "Pole", "Road_marking", "Road", "Pavement", "Tree", "SignSymbol", "Fence", "Car", "Pedestrian", "Bicyclist"]

    ax.set_xticks(np.arange(confusion_matrix.shape[0]) + 0.5, minor=False)
    ax.set_yticks(np.arange(confusion_matrix.shape[1]) + 0.5, minor=False)

    ax.invert_yaxis()
    ax.xaxis.tick_top()

    ax.set_xticklabels(labels, minor=False)
    ax.set_yticklabels(labels, minor=False)
    plt.title("Confusion Matrix acc={}".format(acc))
    plt.savefig("confusion_matrix.png")
    print("saved confusion matrix")

実行

トレーニングを開始するためには、config.pyのbatchsizeなどを設定した後に

python3 train.py

としてください。すると、下の画像のようにepoch数, train loss, valid loss, 経過時間が順にターミナル上に出力されていきます。

10epoch毎にloss_curve.pngが更新されていきます。30epochで既に過学習が起きていることが分かります。@uni-3さんの記事にある結果よりも過学習が遅く始まるのは、dropoutの効果と思われます。

テストを行うためには

python3 test.py

を実行します。dataset/test内の画像を1枚ずつ訓練済みモデルに渡して、出てきた2次元配列を着色したものをpredicted_imgsに保存します。その後、混合行列を作ります。

画像を見てわかるように非常に汎化性能は低いです。もっと沢山訓練データが必要なようです。

最後に

SegNetを実装してみましたが、見直してみるといろいろと反省する点がでてきました。

・classが使いこなせてない。loaderとか普通にクラスにするべきでは？
・みんなが嫌いなglobal変数を使いまくってる。
・PILとcv2が混ざっていて気持ち悪い。

ところで、ここで使われているCamVidデータセットの詳細ってどこに書いてあるのでしょうか。クラス数とか、class_weightingとかは@uni-3さんの記事から持ってきたのですが、1次ソースが見当たらないです。誰か知っていたら教えてください。

追記 U-Netについて

SegNetのエンコード・デコード間にスキップ結合を含めたU-Netも実装してみました。Unetについての詳しい説明は@tktktks10を参照してください。ネットワークは原論文を参考にして以下のようにしました。

network.py

class UNet(Chain):

    def __init__(self, n_class=n_class):
        super().__init__()
        with self.init_scope():
            self.n_class = n_class

            self.enco1_1 = L.Convolution2D(None,  64, ksize=3, stride=1, pad=1)
            self.enco1_2 = L.Convolution2D(None,  64, ksize=3, stride=1, pad=1)

            self.enco2_1 = L.Convolution2D(None, 128, ksize=3, stride=1, pad=1)
            self.enco2_2 = L.Convolution2D(None, 128, ksize=3, stride=1, pad=1)

            self.enco3_1 = L.Convolution2D(None, 256, ksize=3, stride=1, pad=1)
            self.enco3_2 = L.Convolution2D(None, 256, ksize=3, stride=1, pad=1)

            self.enco4_1 = L.Convolution2D(None, 512, ksize=3, stride=1, pad=1)
            self.enco4_2 = L.Convolution2D(None, 512, ksize=3, stride=1, pad=1)

            self.enco5_1 = L.Convolution2D(None,1012, ksize=3, stride=1, pad=1)

            self.deco6_1 = L.Convolution2D(None,1012, ksize=3, stride=1, pad=1)
            self.deco6_2 = L.Convolution2D(None, 512, ksize=3, stride=1, pad=1)

            self.deco7_1 = L.Convolution2D(None, 512, ksize=3, stride=1, pad=1)
            self.deco7_2 = L.Convolution2D(None, 256, ksize=3, stride=1, pad=1)

            self.deco8_1 = L.Convolution2D(None, 256, ksize=3, stride=1, pad=1)
            self.deco8_2 = L.Convolution2D(None, 128, ksize=3, stride=1, pad=1)

            self.deco9_1 = L.Convolution2D(None, 128, ksize=3, stride=1, pad=1)
            self.deco9_2 = L.Convolution2D(None,  64, ksize=3, stride=1, pad=1)
            self.deco9_3 = L.Convolution2D(None,  64, ksize=3, stride=1, pad=1)

            self.final_layer = L.Convolution2D(None, n_class, ksize=1)

            self.bn1_1 = L.BatchNormalization(  64)
            self.bn1_2 = L.BatchNormalization(  64)

            self.bn2_1 = L.BatchNormalization( 128)
            self.bn2_2 = L.BatchNormalization( 128)

            self.bn3_1 = L.BatchNormalization( 256)
            self.bn3_2 = L.BatchNormalization( 256)

            self.bn4_1 = L.BatchNormalization( 512)
            self.bn4_2 = L.BatchNormalization( 512)

            self.bn5_1 = L.BatchNormalization(1012)

            self.bn6_1 = L.BatchNormalization(1012)
            self.bn6_2 = L.BatchNormalization( 512)

            self.bn7_1 = L.BatchNormalization( 512)
            self.bn7_2 = L.BatchNormalization( 256)

            self.bn8_1 = L.BatchNormalization( 256)
            self.bn8_2 = L.BatchNormalization( 128)

            self.bn9_1 = L.BatchNormalization( 128)
            self.bn9_2 = L.BatchNormalization(  64)
            self.bn9_3 = L.BatchNormalization(  64)

    def __call__(self, x): #x = (batchsize, 3, 360, 480)
        if LRN:
            x = F.local_response_normalization(x) #Needed for preventing from overfitting

        h1_1 = F.relu(self.bn1_1(self.enco1_1(x)))
        h1_2 = F.relu(self.bn1_2(self.enco1_2(h1_1)))

        pool1 = F.max_pooling_2d(h1_2, 2, stride=2, return_indices=False) #(batchsize,  64, 180, 240)

        h2_1 = F.relu(self.bn2_1(self.enco2_1(pool1)))
        h2_2 = F.relu(self.bn2_2(self.enco2_2(h2_1)))
        pool2 = F.max_pooling_2d(h2_2, 2, stride=2, return_indices=False) #(batchsize, 128,  90, 120) 

        h3_1 = F.relu(self.bn3_1(self.enco3_1(pool2)))
        h3_2 = F.relu(self.bn3_2(self.enco3_2(h3_1)))
        pool3 = F.max_pooling_2d(h3_2, 2, stride=2, return_indices=False) #(batchsize, 256,  45,  60) 

        h4_1 = F.relu(self.bn4_1(self.enco4_1(pool3)))
        h4_2 = F.relu(self.bn4_2(self.enco4_2(h4_1)))
        pool4 = F.max_pooling_2d(h4_2, 2, stride=2, return_indices=False) #(batchsize, 256,  23,  30) 

        h5_1 = F.relu(self.bn5_1(self.enco5_1(pool4)))

        up5 = F.unpooling_2d(h5_1, ksize=2, stride=2, outsize=(pool3.shape[2], pool3.shape[3]))
        h6_1 = F.relu(self.bn6_1(self.deco6_1(F.concat((up5, h4_2)))))
        h6_2 = F.relu(self.bn6_2(self.deco6_2(h6_1)))

        up6 = F.unpooling_2d(h6_2, ksize=2, stride=2, outsize=(pool2.shape[2], pool2.shape[3]))
        h7_1 = F.relu(self.bn7_1(self.deco7_1(F.concat((up6, h3_2)))))
        h7_2 = F.relu(self.bn7_2(self.deco7_2(h7_1)))

        up7 = F.unpooling_2d(h7_2, ksize=2, stride=2, outsize=(pool1.shape[2], pool1.shape[3]))
        h8_1 = F.relu(self.bn8_1(self.deco8_1(F.concat((up7, h2_2)))))
        h8_2 = F.relu(self.bn8_2(self.deco8_2(h8_1)))

        up8 = F.unpooling_2d(h8_2, ksize=2, stride=2, outsize=(x.shape[2], x.shape[3])) #x = (batchsize, 128, 360, 480)
        h9_1 = F.relu(self.bn9_1(self.deco9_1(F.concat((up8, h1_2)))))
        h9_2 = F.relu(self.bn9_2(self.deco9_2(h9_1)))
        h9_3 = F.relu(self.bn9_3(self.deco9_3(h9_2)))

        h = self.final_layer(h9_3)

        return h

テスト結果のSegNetとU-Netとの比較は以下の通り。U-Netの方がSegNetよりも解像度の高い結果が得られている。

lossの変化は以下の通り。SegNetの場合に比べて過学習が起きるのが遅いのがわかる。