NiftyNetで3D U-Netを実装する

前回の記事でNiftyNetで3次元セグメンテーションを実装しました。
https://qiita.com/takeajioka/items/2a128639dde4c5c97b67
前回はDenseVnetというネットワークを用いましたが、今回はより有名な3D U-Netで実装します。
データセットは前回と同じKaggleのElectron Microscopy 3D Segmentationを用いました。

3D U-Netの概要

U-Netは比較的少ない教師データで精度よく学習できると言われています。
https://lmb.informatik.uni-freiburg.de/Publications/2016/CABR16/
論文を読むとわかる通り、inputデータのサイズとoutputデータのサイズが一致していないということがわかります。
niftyNetではこのネットワーク構造を忠実に再現して使うため、window sizeを決めるときに注意しなければなりません。

実装する

3D U-Netのconfigは以下のようになりました。
参考にしたコードなど：https://cmiclab.cs.ucl.ac.uk/CMIC/NiftyNet/blob/dev/config/unet_config.ini
https://github.com/NifTK/NiftyNet/issues/17
https://stackoverflow.com/questions/46831160/niftynet-indices-are-out-of-bounds-error

############################ input configuration sections
[EM]
path_to_search = ./data/EM3D/
filename_contains = input
spatial_window_size = (96, 96, 96)

[label]
path_to_search = ./data/EM3D/
filename_contains = groundtruth
spatial_window_size = (8, 8, 8)

############################## system configuration sections
[SYSTEM]
cuda_devices = ""
num_threads = 1
num_gpus = 1
model_dir = models/EM3D/

[NETWORK]
name = unet
activation_function = prelu
batch_size = 1
decay = 0
reg_type = L2
volume_padding_size = 44
queue_length = 1
window_sampling = uniform

[TRAINING]
optimiser = adam
sample_per_volume = 1
lr = 0.01
loss_type = CrossEntropy
starting_iter = 0
save_every_n = 100
max_iter = 100

[INFERENCE]
border = (44, 44, 44)
save_seg_dir = ./segmentation_output/

############################ custom configuration sections
[SEGMENTATION]
image = EM
label = label
label_normalisation = False
output_prob = False
num_classes = 2

重要な点は、
1.spatial_window_sizeをoutputではinput - 88にする。
2.Volume paddingを44に設定する。
3.borderも44に設定する。
という点です。
この設定が3D U-Netには必要で、これをしないとエラーになりました。
パラメータの意味は以下に説明があります。
https://niftynet.readthedocs.io/en/dev/window_sizes.html

またWindows sizeが小さいですが、GPUメモリの都合上これが上限です。
Windows sizeをこれ以上小さくすることも出来ないので、３D U-Netには8GB以上のGPUが必要だとわかりました。

ちなみに通常の2次元でのU-NetもNiftyNetで実装することが出来て、configファイルは以下が参考になります。
https://miccai-sb.github.io/materials/U-Net_Demo.html

実行コマンドは以下です。

net_segment train -c <configファイルのパス>
net_segment inference -c <configファイルのパス>

inferenceの方はものすごく時間がかかりました。(spatial_window_sizeが(8, 8, 8)で小さなタイルで全体を処理をするため。)

感想

volume_paddingをしないといけない点がとても気になります。
対策として、niftynetパッケージのpyファイルを編集することでvolume_paddingをしないようにすることも一応は可能のようです。
また、V-netなどの他のネットワークではvolume_paddingは必要ないのでそっちを使う手もあります。

niftynet以外では、kerasなどでも３D U-Netを実装したものはもちろんあるようです。
(https://github.com/ellisdg/3DUnetCNN)
データ元のKaggleページのKernelsにはconxを用いたU-Netの可視化方法も解説されているので参考にしたいです。
またimageJプラグインのU-Netでは電子顕微鏡モデルのfine-tuningを行うことができるのでそれを利用する方法もよいと思います。(プラグインの利用方法は過去の記事を参照https://qiita.com/takeajioka/items/ad13fa0ed19aea48ec2b)