def __init__(self): super().__init__() self.init_block = nn.Sequential( nn.Conv2d(1, 32, kernel_size=3, padding=1), layers.PreActivation2d(32, 32, kernel_size=3, padding=1), ) # don't need upsampling here because it will be done in `merge_branches` upsample = nn.Sequential # same as nn.Identity, but supported by older versions downsample = partial(nn.MaxPool2d, kernel_size=2) self.midline_head = layers.InterpolateToInput( nn.Sequential( downsample(), # unet layers.FPN(layers.ResBlock2d, downsample=downsample, upsample=upsample, merge=merge_branches, structure=STRUCTURE, kernel_size=3, padding=1), # final logits layers.PreActivation2d(32, 1, kernel_size=1), ), mode='bilinear') self.limits_head = layers.InterpolateToInput( nn.Sequential( layers.ResBlock2d(32, 32, kernel_size=3, padding=1), downsample(), layers.ResBlock2d(32, 32, kernel_size=3, padding=1), # 2D feature map to 1D feature map nn.AdaptiveMaxPool2d((None, 1)), layers.Reshape('0', '1', -1), nn.Conv1d(32, 8, kernel_size=3, padding=1), nn.ReLU(), nn.Conv1d(8, 1, kernel_size=3, padding=1)), mode='linear', axes=0)
def build_encoder(nums_conv, channels, n_features, n_narrow, kernel_sizes, padding, stride, poolings): ''' :param nums_conv: int, number of convolutional blocks :param channels: List of ints, the length should be divided by 2. Provides number of channels for each convolution. :param n_features: Shape of output after convolutions :param n_narrow: Shape to inner representation :param kernel_sizes: Size of kernel in convolutions :param padding: Padding size in convolutions :param stride: Stride size in convolutions :param poolings: The list of poolings after blocks :return: Encoder part: CNN + fc ''' encoder = nn.Sequential( convs(nums_conv, channels, kernel_sizes, [padding] * nums_conv, [stride] * nums_conv, poolings, nn.Sequential(), conv_module=nn.Conv2d), layers.Reshape('0', -1), nn.Linear(n_features, n_narrow)) return encoder
def __init__(self, n_classes=2, in_channels=3, feature_scale=4, is_batchnorm=True, up_sample_mode='transpose_conv', beta=10): super(Ours_CARNet, self).__init__() self.up_sample_mode = up_sample_mode self.in_channels = in_channels self.is_batchnorm = is_batchnorm self.feature_scale = feature_scale self.beta = beta filters = [64, 128, 256, 512, 1024] filters = [int(x / self.feature_scale) for x in filters] downsample = partial(nn.MaxPool2d, kernel_size=2) self.limits_head = layers.InterpolateToInput( nn.Sequential( layers.ResBlock2d(16, 32, kernel_size=3, padding=1), downsample(), layers.ResBlock2d(32, 32, kernel_size=3, padding=1), # 2D feature map to 1D feature map nn.AdaptiveMaxPool2d((None, 1)), layers.Reshape('0', '1', -1), nn.Conv1d(32, 8, kernel_size=3, padding=1), nn.ReLU(), nn.Conv1d(8, 1, kernel_size=3, padding=1)), mode='linear', axes=0) # downsampling self.conv1 = unetConv2(self.in_channels, filters[0], self.is_batchnorm) self.maxpool1 = nn.MaxPool2d(kernel_size=2) self.conv2 = unetConv2(filters[0], filters[1], self.is_batchnorm) self.maxpool2 = nn.MaxPool2d(kernel_size=2) self.conv3 = unetConv2(filters[1], filters[2], self.is_batchnorm) self.maxpool3 = nn.MaxPool2d(kernel_size=2) self.conv4 = unetConv2(filters[2], filters[3], self.is_batchnorm) self.maxpool4 = nn.MaxPool2d(kernel_size=2) self.center = unetConv2(filters[3], filters[4], self.is_batchnorm) # upsampling self.up_concat4 = unetUp(filters[4], filters[3], self.up_sample_mode) self.up_concat3 = unetUp(filters[3], filters[2], self.up_sample_mode) self.up_concat2 = unetUp(filters[2], filters[1], self.up_sample_mode) self.up_concat1 = unetUp(filters[1], filters[0], self.up_sample_mode) # final conv (without any concat) self.final_1 = nn.Conv2d(filters[0], 1, 1) self.sigmoid = nn.Sigmoid() # ===== refine net ===== # self.laternal_1 = nn.Conv2d(filters[0], filters[0], 1) self.laternal_2 = nn.Conv2d(filters[1], filters[0], 1) self.laternal_3 = nn.Conv2d(filters[2], filters[0], 1) self.laternal_4 = nn.Conv2d(filters[3], filters[0], 1) self.laternal_5 = nn.Conv2d(filters[4], filters[0], 1) self.refine_conv_1 = SELayer(filters[0]) self.refine_conv_2 = nn.Sequential(*[ unetConv2(filters[0], filters[0], self.is_batchnorm),\ SELayer(filters[0]) ]) self.refine_conv_3 = nn.Sequential(*[ unetConv2(filters[0], filters[0], self.is_batchnorm),\ unetConv2(filters[0], filters[0], self.is_batchnorm),\ SELayer(filters[0]) ]) self.refine_conv_4 = nn.Sequential(*[ unetConv2(filters[0], filters[0], self.is_batchnorm),\ unetConv2(filters[0], filters[0], self.is_batchnorm),\ unetConv2(filters[0], filters[0], self.is_batchnorm),\ SELayer(filters[0]) ]) self.refine_conv_5 = nn.Sequential(*[ unetConv2(filters[0], filters[0], self.is_batchnorm),\ unetConv2(filters[0], filters[0], self.is_batchnorm),\ unetConv2(filters[0], filters[0], self.is_batchnorm),\ unetConv2(filters[0], filters[0], self.is_batchnorm),\ SELayer(filters[0]) ]) self.refine_final = unetConv2(filters[0] * 5, filters[0], self.is_batchnorm) self.final_2 = nn.Conv2d(filters[0], 1, 1) # ===== regression_based_refine ===== # self.downsample_a = nn.Sequential( nn.Conv2d(1, 32, 3, stride=2, padding=1), nn.BatchNorm2d(32)) self.res_blocks_a = BasicBlock(inplanes=1, planes=32, stride=2, downsample=self.downsample_a) self.downsample_b = nn.Sequential( nn.Conv2d(32, 64, 3, stride=2, padding=1), nn.BatchNorm2d(64)) self.res_blocks_b = BasicBlock(inplanes=32, planes=64, stride=2, downsample=self.downsample_b) self.downsample_c = nn.Sequential( nn.Conv2d(64, 128, 3, stride=2, padding=1), nn.BatchNorm2d(128)) self.res_blocks_c = BasicBlock(inplanes=64, planes=128, stride=2, downsample=self.downsample_c) self.softargmax = SoftArgmax(beta=self.beta) self.res_blocks_mid = nn.Sequential( *[self.res_blocks_a, self.res_blocks_b, self.res_blocks_c]) self.avgpool = nn.AdaptiveAvgPool2d((4, 4)) self.regress_mid = nn.Sequential(nn.Linear(128 * 4 * 4, 600), nn.Linear(600, 400)) # initialise weights for m in self.modules(): if isinstance(m, nn.Conv2d): init_weights(m, init_type='kaiming') elif isinstance(m, nn.BatchNorm2d): init_weights(m, init_type='kaiming')