def __init__(self, in_size, gate_size, inter_size, nonlocal_mode,
sub_sample_factor):
super(MultiAttentionBlock, self).__init__()
self.gate_block_1 = GridAttentionBlock3D(
in_channels=in_size,
gating_channels=gate_size,
inter_channels=inter_size,
mode=nonlocal_mode,
sub_sample_factor=sub_sample_factor)
self.gate_block_2 = GridAttentionBlock3D(
in_channels=in_size,
gating_channels=gate_size,
inter_channels=inter_size,
mode=nonlocal_mode,
sub_sample_factor=sub_sample_factor)
self.combine_gates = nn.Sequential(
nn.Conv3d(in_size * 2, in_size, kernel_size=1,
stride=1, padding=0), nn.BatchNorm3d(in_size),
nn.ReLU(inplace=True))
# initialise the blocks
for m in self.children():
if m.__class__.__name__.find('GridAttentionBlock3D') != -1:
continue
init_weights(m, init_type='kaiming')
def __init__(self, feature_scale=4, n_classes=21, is_deconv=True, in_channels=3,
nonlocal_mode='concatenation', attention_dsample=(2,2,2), is_batchnorm=True):
super(unet_grid_attention_3D, self).__init__()
self.is_deconv = is_deconv
self.in_channels = in_channels
self.is_batchnorm = is_batchnorm
self.feature_scale = feature_scale
filters = [64, 128, 256, 512, 1024]
filters = [int(x / self.feature_scale) for x in filters]
# downsampling
self.conv1 = UnetConv3(self.in_channels, filters[0], self.is_batchnorm)
self.maxpool1 = nn.MaxPool3d(kernel_size=(2, 2, 1))
self.conv2 = UnetConv3(filters[0], filters[1], self.is_batchnorm)
self.maxpool2 = nn.MaxPool3d(kernel_size=(2, 2, 1))
self.conv3 = UnetConv3(filters[1], filters[2], self.is_batchnorm)
self.maxpool3 = nn.MaxPool3d(kernel_size=(2, 2, 1))
self.conv4 = UnetConv3(filters[2], filters[3], self.is_batchnorm)
self.maxpool4 = nn.MaxPool3d(kernel_size=(2, 2, 1))
self.center = UnetConv3(filters[3], filters[4], self.is_batchnorm)
self.gating = UnetGridGatingSignal3(filters[4], filters[3], kernel_size=(1, 1, 1), is_batchnorm=self.is_batchnorm)
# attention blocks
self.attentionblock2 = GridAttentionBlock3D(in_channels=filters[1], gating_channels=filters[3],
inter_channels=filters[1], sub_sample_factor=attention_dsample, mode=nonlocal_mode)
self.attentionblock3 = GridAttentionBlock3D(in_channels=filters[2], gating_channels=filters[3],
inter_channels=filters[2], sub_sample_factor=attention_dsample, mode=nonlocal_mode)
self.attentionblock4 = GridAttentionBlock3D(in_channels=filters[3], gating_channels=filters[3],
inter_channels=filters[3], sub_sample_factor=attention_dsample, mode=nonlocal_mode)
# upsampling
self.up_concat4 = UnetUp3(filters[4], filters[3], self.is_deconv, self.is_batchnorm)
self.up_concat3 = UnetUp3(filters[3], filters[2], self.is_deconv, self.is_batchnorm)
self.up_concat2 = UnetUp3(filters[2], filters[1], self.is_deconv, self.is_batchnorm)
self.up_concat1 = UnetUp3(filters[1], filters[0], self.is_deconv, self.is_batchnorm)
# final conv (without any concat)
self.final = nn.Conv3d(filters[0], n_classes, 1)
# initialise weights
for m in self.modules():
if isinstance(m, nn.Conv3d):
init_weights(m, init_type='kaiming')
elif isinstance(m, nn.BatchNorm3d):
init_weights(m, init_type='kaiming')