class UNet(nn.Module):
def __init__(self, depth=5, root_filter=64):
super(UNet, self).__init__()
self.depth = depth
self.root_filter = root_filter
print("depth: {}, root filter: {}".format(self.depth, self.root_filter))
def forward(self, x):
long_connection_store = {}
input_channels = 1
# Down sampling
print("DOWN SAMPLING")
for i in range(self.depth):
out_channel = 2**i * self.root_filter
print(out_channel)
x = F.relu(nn.Conv2d(input_channels, out_channel, 3, padding=1)(x))
x = F.relu(nn.Conv2d(out_channel, out_channel, 3, padding=1)(x))
if i<self.depth-1:
long_connection_store[str(i)] = x
x = F.max_pool2d(x, kernel_size=2)
input_channels = out_channel
print("x.shape: {}".format(x.shape))
print("UP SAMPLING")
for i in range(self.depth - 2, -1, -1):
out_channel = 2**(i) * self.root_filter
print("input_channels: {}".format(input_channels))
print("out_channel: {}".format(out_channel))
# long connection from down sampling path.
long_connection = long_connection_store[str(i)]
print("long_connection shape: {}".format(long_connection.shape))
up1 = nn.Upsample(scale_factor=2)(x)
print("up1 shape: {}".format(up1.shape))
up_conv1 = F.relu(nn.Conv2d(input_channels, out_channel, 3, padding=1)(up1))
print("upsamplingConv : {}".format(up_conv1.shape))
up_conc = torch.cat((up_conv1, long_connection), dim=1)
print("up_conc : {}\n".format(up_conc.shape))
# Convolutions
up_conv1 = F.relu(nn.Conv2d(input_channels, out_channel, 3, padding=1)(up_conc))
up_conv2 = F.relu(nn.Conv2d(out_channel, out_channel, 3, padding=1)(up_conv1))
x = up_conv2
input_channels = out_channel
print("FINAL shape: {}".format(x.shape))
x = F.sigmoid(nn.Conv2d(input_channels, 1, 1, padding=0)(x))
print("FINAL shape: {}".format(x.shape))
return x
batch_size = 5
channels = 1
height = 256
width = 256
image = torch.randn(batch_size, channels, height, width)
unet = UNet(5, 64)
output = unet(image)
If I run the above code, it works but when I do unet.parameters()
I get empty list.
Is it always necessary to define the layers with learnable parameters in init function? That looks like lot of work if the model is very deep like in case of UNet. Do I have to define all Conv layers of downsampling and upsampling path in init function. If I do that then I can’t initialize the model dynamically with depth or filter size parameters which I can do very easily in Keras.
Is there a way to define the model without initiating each and every Conv layer in init function.