#

ユーネット

これは、論文「U-Net: 生物医学画像セグメンテーションのための畳み込みネットワーク」のU-Netモデルの実装です。

U-Netは収縮経路と拡張経路で構成されています。収縮経路は一連の畳み込み層とプーリング層であり、特徴マップの解像度は徐々に低下します。エクスパンシブパスとは、フィーチャマップの解像度が徐々に上がっていく一連のアップサンプリングレイヤーと畳み込みレイヤーのことです

。

拡張パスの各ステップで、縮小パスからの対応するフィーチャマップが現在のフィーチャマップと連結されます。

U-Net diagram from paper

これは、CarvanaデータセットでU-Netをトレーニングする実験のトレーニングコードです。

27import torch
28import torchvision.transforms.functional
29from torch import nn

#

2 $3 \times 3$ つのコンボリューションレイヤー

収縮経路と膨張経路の各ステップには、 $3 \times 3$ 2つの畳み込み層があり、その後にReLU活性化が続きます。

$0$ U-Netの論文ではパディングを使用していましたが、 $1$ 最終的なフィーチャマップがトリミングされないようにパディングを使用しています。

32class DoubleConvolution(nn.Module):

#

in_channels は入力チャンネル数
out_channels は出力チャンネル数

43    def __init__(self, in_channels: int, out_channels: int):

#

48        super().__init__()

#

$3 \times 3$ 最初の畳み込み層

51        self.first = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
52        self.act1 = nn.ReLU()

#

2 $3 \times 3$ 番目の畳み込み層

54        self.second = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
55        self.act2 = nn.ReLU()

#

57    def forward(self, x: torch.Tensor):

#

2 つのコンボリューションレイヤーとアクティベーションを適用します。

59        x = self.first(x)
60        x = self.act1(x)
61        x = self.second(x)
62        return self.act2(x)

#

ダウンサンプル

収縮パスの各ステップは、 $2 \times 2$ 最大プーリングレイヤーを使用して特徴マップをダウンサンプリングします。

65class DownSample(nn.Module):

#

73    def __init__(self):
74        super().__init__()

#

最大プーリング層

76        self.pool = nn.MaxPool2d(2)

#

78    def forward(self, x: torch.Tensor):
79        return self.pool(x)

#

アップサンプル

広大な経路の各ステップは、アップコンボリューションで特徴マップをアップサンプリングします。 $2 \times 2$

82class UpSample(nn.Module):

#

89    def __init__(self, in_channels: int, out_channels: int):
90        super().__init__()

#

アップコンボリューション

93        self.up = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2)

#

95    def forward(self, x: torch.Tensor):
96        return self.up(x)

#

フィーチャマップのトリミングと連結

拡張パスの各ステップで、縮小パスからの対応するフィーチャマップが現在のフィーチャマップと連結されます。

99class CropAndConcat(nn.Module):

#

x エクスパンシブパス内の現在のフィーチャマップ
contracting_x 契約経路からの対応する機能マップ

106    def forward(self, x: torch.Tensor, contracting_x: torch.Tensor):

#

フィーチャマップを縮小パスから現在のフィーチャマップのサイズにトリミングします

113        contracting_x = torchvision.transforms.functional.center_crop(contracting_x, [x.shape[2], x.shape[3]])

#

フィーチャマップを連結する

115        x = torch.cat([x, contracting_x], dim=1)

#

117        return x

#