def __init__(self, in_channels, out_channels, final_sigmoid, basic_module, f_maps=64, layer_order='gcr', num_groups=8, num_levels=4, is_segmentation=True, conv_kernel_size=3, pool_kernel_size=2, conv_padding=1, **kwargs): super(Abstract3DUNet, self).__init__() if isinstance(f_maps, int): f_maps = number_of_features_per_level(f_maps, num_levels=num_levels) assert isinstance(f_maps, list) or isinstance(f_maps, tuple) assert len(f_maps) > 1, "Required at least 2 levels in the U-Net" # create encoder path self.encoders = create_encoders(in_channels, f_maps, basic_module, conv_kernel_size, conv_padding, layer_order, num_groups, pool_kernel_size) # create decoder path self.decoders = create_decoders(f_maps, basic_module, conv_kernel_size, conv_padding, layer_order, num_groups, upsample=True) # in the last layer a 1×1 convolution reduces the number of output # channels to the number of labels self.final_conv = nn.Conv3d(f_maps[0], out_channels, 1) if is_segmentation: # semantic segmentation problem if final_sigmoid: self.final_activation = nn.Sigmoid() else: self.final_activation = nn.Softmax(dim=1) else: # regression problem self.final_activation = None
def __init__(self, in_channels, out_channels, final_sigmoid, basic_module, f_maps=64, layer_order='gcr', num_groups=8, num_levels=4, is_segmentation=True, testing=False, **kwargs): super(Abstract3DUNet, self).__init__() self.testing = testing if isinstance(f_maps, int): f_maps = number_of_features_per_level(f_maps, num_levels=num_levels) # create encoder path consisting of Encoder modules. Depth of the encoder is equal to `len(f_maps)` encoders = [] for i, out_feature_num in enumerate(f_maps): if i == 0: encoder = Encoder(in_channels, out_feature_num, apply_pooling=False, basic_module=basic_module, conv_layer_order=layer_order, num_groups=num_groups) else: # TODO: adapt for anisotropy in the data, i.e. use proper pooling kernel to make the data isotropic after 1-2 pooling operations # currently pools with a constant kernel: (2, 2, 2) encoder = Encoder(f_maps[i - 1], out_feature_num, basic_module=basic_module, conv_layer_order=layer_order, num_groups=num_groups) encoders.append(encoder) self.encoders = nn.ModuleList(encoders) # create decoder path consisting of the Decoder modules. The length of the decoder is equal to `len(f_maps) - 1` decoders = [] reversed_f_maps = list(reversed(f_maps)) for i in range(len(reversed_f_maps) - 1): if basic_module == DoubleConv: in_feature_num = reversed_f_maps[i] + reversed_f_maps[i + 1] else: in_feature_num = reversed_f_maps[i] out_feature_num = reversed_f_maps[i + 1] # TODO: if non-standard pooling was used, make sure to use correct striding for transpose conv # currently strides with a constant stride: (2, 2, 2) decoder = Decoder(in_feature_num, out_feature_num, basic_module=basic_module, conv_layer_order=layer_order, num_groups=num_groups) decoders.append(decoder) self.decoders = nn.ModuleList(decoders) # in the last layer a 1×1 convolution reduces the number of output # channels to the number of labels self.final_conv = nn.Conv3d(f_maps[0], out_channels, 1) if is_segmentation: # semantic segmentation problem if final_sigmoid: self.final_activation = nn.Sigmoid() else: self.final_activation = nn.Softmax(dim=1) else: # regression problem self.final_activation = None