Pytorch で U-Net(セマンティックセグメンテーション)

おんちゃんは、U-Net の勉強が目的なので、学習を試してみました。

>python train.py -g

を実行したら、

Starting training: Epochs: 5 Batch size: 5 Learning rate: 0.1 Training size: 9668 Validation size: 508 Checkpoints: True CUDA: True Starting epoch 1/5. Traceback (most recent call last): File "C:\workspace-home\python-study\Pytorch-UNet\train.py", line 152, in <module> img_scale=args.scale) File "C:\workspace-home\python-study\Pytorch-UNet\train.py", line 88, in train_net masks_pred = net(imgs) File "C:\app\Anaconda3\envs\pytorch\lib\site-packages\torch\nn\modules\module.py", line 493, in __call__ result = self.forward(*input, **kwargs) File "C:\workspace-home\python-study\Pytorch-UNet\unet\unet_model.py", line 30, in forward x = self.up4(x, x1) File "C:\app\Anaconda3\envs\pytorch\lib\site-packages\torch\nn\modules\module.py", line 493, in __call__ result = self.forward(*input, **kwargs) File "C:\workspace-home\python-study\Pytorch-UNet\unet\unet_parts.py", line 77, in forward x = self.conv(x) File "C:\app\Anaconda3\envs\pytorch\lib\site-packages\torch\nn\modules\module.py", line 493, in __call__ result = self.forward(*input, **kwargs) File "C:\workspace-home\python-study\Pytorch-UNet\unet\unet_parts.py", line 22, in forward x = self.conv(x) File "C:\app\Anaconda3\envs\pytorch\lib\site-packages\torch\nn\modules\module.py", line 493, in __call__ result = self.forward(*input, **kwargs) File "C:\app\Anaconda3\envs\pytorch\lib\site-packages\torch\nn\modules\container.py", line 92, in forward input = module(input) File "C:\app\Anaconda3\envs\pytorch\lib\site-packages\torch\nn\modules\module.py", line 493, in __call__ result = self.forward(*input, **kwargs) File "C:\app\Anaconda3\envs\pytorch\lib\site-packages\torch\nn\modules\conv.py", line 338, in forward self.padding, self.dilation, self.groups) RuntimeError: CUDA out of memory. Tried to allocate 180.00 MiB (GPU 0; 4.00 GiB total capacity; 2.80 GiB already allocated; 117.05 MiB free; 38.70 MiB cached)

GPUのメモリー不足との事でした。

同じ、不具合の件が出ていたの、参考にしたところ、どうやら、
入力画像のサイズとバッチサイズを小さくすれば良いとの情報が得られました。
https://discuss.pytorch.org/t/gpu-memory-issue-on-unet/42651/4

>python train.py -g -b 3 -s 0.3
で、どうやら、GPU でOK となりました。

画像の適性サイズは、def train_net() の中で、 summary(net,input_size=(3,width,hight)) で確認できます。
def train_net(net, .... optimizer = optim.SGD(net.parameters(), lr=lr, momentum=0.9, weight_decay=0.0005) criterion = nn.BCELoss() from torchsummary import summary ''' 下の入力画像サイズを変えて、GPUメモリー不足にならないか、当たりを付けます。 ''' summary(net,input_size=(3,512,512)) # ここの画像サイズを変えて試す。 for epoch in range(epochs): ...