I am trying to create a 3D unet for medical image segmentation with dilation layers to give me a the required receptive field size i need without making the model too heavy. I am having a lot of trouble fixing the sizes along the encoder and decoder to match. My input is [bs,4,96,96,64] and output [bs,5,96,96,64] i.e each pixel has 5 possible class possibilities. As you see I have my down block and upblock defined and I max pool two times in the encoder part. I have shown the sizes and error after the code below:
The error is in the first upblock as that interpolates a size of 12 while the corrposing block is size 13 in the 3rd dimension. Can someone please help me make it symmetrical or atleast workable?
class UNet_down_block(torch.nn.Module):
def __init__(self, input_channel, output_channel, down_size):
super(UNet_down_block, self).__init__()
self.conv1 = torch.nn.Conv3d(input_channel, output_channel, 3, padding=1,dilation=1)
self.bn1 = torch.nn.BatchNorm3d(output_channel)
self.conv2 = torch.nn.Conv3d(output_channel, output_channel, 3, padding=1,dilation=2)
self.bn2 = torch.nn.BatchNorm3d(output_channel)
self.conv3 = torch.nn.Conv3d(output_channel, output_channel, 3, padding=1,dilation=2)
self.bn3 = torch.nn.BatchNorm3d(output_channel)
self.max_pool = torch.nn.MaxPool3d(2, 2)
self.relu = torch.nn.ELU()
self.down_size = down_size
def forward(self, x):
if self.down_size:
x = self.max_pool(x)
x = self.relu(self.bn1(self.conv1(x)))
x = self.relu(self.bn2(self.conv2(x)))
x = self.relu(self.bn3(self.conv3(x)))
return x
class UNet_up_block(torch.nn.Module):
def __init__(self, prev_channel, input_channel, output_channel):
super(UNet_up_block, self).__init__()
# self.up_sampling = torch.nn.functional.interpolate(scale_factor=2, mode='trilinear')
self.conv1 = torch.nn.Conv3d(prev_channel + input_channel, output_channel, 3, padding=1)
self.bn1 = torch.nn.BatchNorm3d(output_channel)
self.conv2 = torch.nn.Conv3d(output_channel, output_channel, 3, padding=1)
self.bn2 = torch.nn.BatchNorm3d(output_channel)
self.conv3 = torch.nn.Conv3d(output_channel, output_channel, 3, padding=1)
self.bn3 = torch.nn.BatchNorm3d(output_channel)
self.relu = torch.nn.ELU()
def forward(self, prev_feature_map, x):
x = torch.nn.functional.interpolate(x,scale_factor=2, mode='trilinear')
x = torch.cat((x, prev_feature_map), dim=1)
x = self.relu(self.bn1(self.conv1(x)))
x = self.relu(self.bn2(self.conv2(x)))
x = self.relu(self.bn3(self.conv3(x)))
return x
class UNet(torch.nn.Module):
def __init__(self):
super(UNet, self).__init__()
self.down_block1 = UNet_down_block(4, 24, False)
self.down_block2 = UNet_down_block(24, 72, True)
self.down_block3 = UNet_down_block(72, 148, True)
self.down_block4 = UNet_down_block(148, 224, False)
self.max_pool = torch.nn.MaxPool3d(2, 2)
self.mid_conv1 = torch.nn.Conv3d(224, 224, 3, padding=1)
self.bn1 = torch.nn.BatchNorm3d(224)
self.mid_conv2 = torch.nn.Conv3d(224, 224, 3, padding=1)
self.bn2 = torch.nn.BatchNorm3d(224)
self.mid_conv3 = torch.nn.Conv3d(224, 224, 3, padding=1)
self.bn3 = torch.nn.BatchNorm3d(224)
self.up_block1 = UNet_up_block(224, 224, 148)
self.up_block2 = UNet_up_block(148, 148, 72)
self.up_block3 = UNet_up_block(72, 72, 24)
self.up_block4 = UNet_up_block(24, 24, 8)
self.last_conv1 = torch.nn.Conv3d(8, 4, 3, padding=1)
self.last_bn = torch.nn.BatchNorm3d(4)
self.last_conv2 = torch.nn.Conv3d(4, 1, 1, padding=0)
self.relu = torch.nn.ELU()
self.last_conv3 = torch.nn.Conv3d(1, 1, 1, padding=0)
self.relu = torch.nn.ELU()
self.conv1f=torch.nn.Conv2d(1, 5, 3,padding=1)
self.conv2f=torch.nn.Conv2d(5, 5, 3,padding=1)
self.conv3f=torch.nn.Conv2d(5, 5, 3,padding=1)
def forward(self, x):
print('input unet',x.size())
self.x1 = self.down_block1(x)
print("Block 1 shape:",self.x1.size())
self.x2 = self.down_block2(self.x1)
if self.x2.size()[2]==49: ###*********************************** ifffff if self.x2.size()[2]==49:
self.x2=self.x2[:,:,1:,1:,:]
print("Block 2 shape:",self.x2.size())
self.x3 = self.down_block3(self.x2)
print("Block 3 shape:",self.x3.size())
self.x4 = self.down_block4(self.x3)
print("Block 4 shape:",self.x4.size())
self.xmid=self.max_pool(self.x4)
self.xmid = self.relu(self.bn1(self.mid_conv1(self.xmid)))
self.xmid = self.relu(self.bn2(self.mid_conv2(self.xmid)))
self.xmid = self.relu(self.bn3(self.mid_conv3(self.xmid)))
print("Block Mid shape:",self.xmid.size())
x = self.up_block1(self.x4, self.xmid)
# print("BlockU 1 shape:",x.size())
x = self.up_block2(self.x3, x)
print("BlockU 2 shape:",x.size())
x = self.up_block3(self.x2, x)
print("BlockU 3 shape:",x.size())
if self.x1.size()[2]==98: ###*********************************** ifffff
self.x1=self.x1[:,:,1:-1,1:-1,:]
# print('chan98',self.x1.size())
x = self.up_block4(self.x1, x)
print("BlockU 4 shape:",x.size())
x = self.relu(self.last_bn(self.last_conv1(x)))
x = self.last_conv2(x) # of size [batch_size,1,h,w,depth] or [bs, modalities(1) ,96 ,96 , 64]
x=x.view(batch_size,1,-1,64)
# x=x.squeeze(1)
# print('input convf',x.size())
conv=self.relu(self.conv1f(x))
conv=self.relu(self.conv2f(conv))
conv=self.conv3f(conv)
try:
conv=conv.view(batch_size,5,96,96,64)
except:
conv=conv.view(batch_size_val,5,96,96,64)
# print('unet output',conv.size())
return(conv)
Here is the the output:
input unet torch.Size([1, 4, 96, 96, 64])
Block 1 shape: torch.Size([1, 24, 92, 92, 60])
Block 2 shape: torch.Size([1, 72, 42, 42, 26])
Block 3 shape: torch.Size([1, 148, 17, 17, 9])
Block 4 shape: torch.Size([1, 224, 13, 13, 5])
Block Mid shape: torch.Size([1, 224, 6, 6, 2])
Error:
x = self.up_block1(self.x4, self.xmid)
111 # print("BlockU 1 shape:",x.size())
112 x = self.up_block2(self.x3, x)
/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
475 result = self._slow_forward(*input, **kwargs)
476 else:
--> 477 result = self.forward(*input, **kwargs)
478 for hook in self._forward_hooks.values():
479 hook_result = hook(self, input, result)
<ipython-input-5-cbcdda025480> in forward(self, prev_feature_map, x)
39 x = torch.nn.functional.interpolate(x,scale_factor=2, mode='trilinear')
40
---> 41 x = torch.cat((x, prev_feature_map), dim=1)
42 x = self.relu(self.bn1(self.conv1(x)))
43 x = self.relu(self.bn2(self.conv2(x)))
RuntimeError: invalid argument 0: Sizes of tensors must match except in dimension 1. Got 12 and 13 in dimension 2 at /opt/conda/conda-bld/pytorch_1535491974311/work/aten/src/TH/generic/THTensorMath.cpp:3616