def __init__(self, in_channels=1, out_channels=2, feature_maps=64, levels=4, skip_connections=True, group_norm=True):
super(UNetDecoder, self).__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.feature_maps = feature_maps
self.levels = levels
self.skip_connections = skip_connections
self.features = nn.Sequential()
for i in range(levels):
# upsampling block
upconv = UNetUpSamplingBlock2D(2**(levels-i) * feature_maps, 2**(levels-i-1) * feature_maps, deconv=True)
self.features.add_module('upconv%d' % (i + 1), upconv)
# convolutional block
if skip_connections:
conv_block = UNetConvBlock2D(2**(levels-i) * feature_maps, 2**(levels-i-1) * feature_maps, group_norm=group_norm)
else:
conv_block = UNetConvBlock2D(2**(levels-i-1) * feature_maps, 2**(levels-i-1) * feature_maps, group_norm=group_norm)
self.features.add_module('convblock%d' % (i + 1), conv_block)
# output layer
self.output = nn.Conv2d(feature_maps, out_channels, kernel_size=1)
def __init__(self, in_channels=1, feature_maps=64, levels=4, group_norm=True):
super(UNetEncoder, self).__init__()
self.in_channels = in_channels
self.feature_maps = feature_maps
self.levels = levels
self.features = nn.Sequential()
in_features = in_channels
for i in range(levels):
out_features = (2**i) * feature_maps
# convolutional block
conv_block = UNetConvBlock2D(in_features, out_features, group_norm=group_norm)
self.features.add_module('convblock%d' % (i + 1), conv_block)
# pooling
pool = nn.MaxPool2d(kernel_size=2, stride=2)
self.features.add_module('pool%d' % (i + 1), pool)
# input features for next block
in_features = out_features
# center (lowest) block
self.center_conv = UNetConvBlock2D(2**(levels-1) * feature_maps, 2**levels * feature_maps * 2, group_norm=group_norm)