def __init__(self, img_ch=1, output_ch=1):
super(AttU_Net, self).__init__()
n1 = 16
filters = [n1, n1 * 2, n1 * 4, n1 * 8, n1 * 16]
self.Maxpool1 = nn.MaxPool3d(kernel_size=2, stride=2)
self.Maxpool2 = nn.MaxPool3d(kernel_size=2, stride=2)
self.Maxpool3 = nn.MaxPool3d(kernel_size=2, stride=2)
self.Maxpool4 = nn.MaxPool3d(kernel_size=2, stride=2)
self.Conv1 = conv_block2(img_ch, filters[0])
self.Conv2 = conv_block(filters[0], filters[1])
self.Conv3 = conv_block(filters[1], filters[2])
self.Conv4 = conv_block(filters[2], filters[3])
self.Conv5 = conv_block(filters[3], filters[4])
self.do1 = nn.Dropout3d(0.5)
self.Up5 = up_conv(filters[4], filters[3])
self.Att5 = Attention_block(F_g=filters[3],
F_l=filters[3],
F_int=filters[2])
self.Up_conv5 = conv_block(filters[4], filters[3])
self.Up4 = up_conv(filters[3], filters[2])
self.Att4 = Attention_block(F_g=filters[2],
F_l=filters[2],
F_int=filters[1])
self.Up_conv4 = conv_block(filters[3], filters[2])
self.Up3 = up_conv(filters[2], filters[1])
self.Att3 = Attention_block(F_g=filters[1],
F_l=filters[1],
F_int=filters[0])
self.Up_conv3 = conv_block(filters[2], filters[1])
self.Up2 = up_conv(filters[1], filters[0])
self.Att2 = Attention_block(F_g=filters[0], F_l=filters[0], F_int=32)
self.Up_conv2 = conv_block(filters[1], filters[0])
self.Conv = nn.Conv3d(filters[0],
output_ch,
kernel_size=1,
stride=1,
padding=0)
self.extra_feature1 = nn.Sequential(nn.AdaptiveAvgPool3d((1, 1, 1)))
self.extra_feature2 = nn.Sequential(nn.AdaptiveAvgPool3d((1, 1, 1)))
self.extra_feature3 = nn.Sequential(nn.AdaptiveAvgPool3d((1, 1, 1)))
self.extra_feature4 = nn.Sequential(nn.AdaptiveAvgPool3d((1, 1, 1)))
self.extra_feature5 = nn.Sequential(nn.AdaptiveAvgPool3d((1, 1, 1)))
self.extra_feature6 = nn.Sequential(nn.AdaptiveAvgPool3d((1, 1, 1)))
self.extra_feature7 = nn.Sequential(nn.AdaptiveAvgPool3d((1, 1, 1)))
self.extra_feature8 = nn.Sequential(nn.AdaptiveAvgPool3d((1, 1, 1)))
self.fc1 = nn.Sequential(nn.Linear(688, 256), nn.Dropout(0.5))
self.fc2 = nn.Sequential(nn.Linear(256, 2), nn.LogSoftmax())
rates = (1, 6, 12, 18)
self.aspp1 = ASPP_module(16, 16, rate=rates[0])
self.aspp2 = ASPP_module(16, 16, rate=rates[1])
self.aspp3 = ASPP_module(16, 16, rate=rates[2])
self.aspp4 = ASPP_module(16, 16, rate=rates[3])
self.aspp_conv = nn.Conv3d(64, 64, 1)
self.aspp_gn = nn.GroupNorm(32, 64)
self.Out = nn.Conv3d(64, output_ch, kernel_size=1, stride=1, padding=0)
def __init__(self,doRate):
super(Dropout3D_, self).__init__()
self.do = nn.Dropout3d(doRate)
def __init__(self, nChans=[16,1], kernel_size=3):
super(MRConvNet, self).__init__()
self.conv1 = nn.Conv3d(1, nChans[0], kernel_size, padding=1)
self.bnorm = nn.BatchNorm3d(nChans[0])
self.drop1 = nn.Dropout3d(p=0.3)
self.conv2 = nn.Conv3d(nChans[0], nChans[1], kernel_size, padding=1)
def __init__(self, column_units):
super(Model, self).__init__()
self.block1 = nn.Sequential(
nn.Conv3d(3,
32,
kernel_size=(3, 5, 5),
stride=(1, 2, 2),
dilation=(1, 1, 1),
padding=(1, 2, 2)),
nn.BatchNorm3d(32),
nn.ReLU(inplace=True),
nn.Dropout3d(p=0.2),
)
self.block2 = nn.Sequential(
nn.Conv3d(32,
64,
kernel_size=(3, 3, 3),
stride=1,
dilation=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(64),
nn.ReLU(inplace=True),
nn.Conv3d(64,
128,
kernel_size=(3, 3, 3),
stride=(1, 2, 2),
dilation=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(128),
nn.ReLU(inplace=True),
nn.AvgPool3d(kernel_size=(3, 1, 1),
stride=(2, 1, 1),
padding=(1, 0, 0)),
nn.Dropout3d(p=0.2),
)
self.block3 = nn.Sequential(
nn.Conv3d(128,
128,
kernel_size=(3, 3, 3),
stride=1,
dilation=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(128),
nn.ReLU(inplace=True),
nn.Conv3d(128,
128,
kernel_size=(3, 3, 3),
stride=1,
dilation=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(128),
nn.ReLU(inplace=True),
nn.Conv3d(128,
256,
kernel_size=(3, 3, 3),
stride=(1, 2, 2),
dilation=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(256),
nn.ReLU(inplace=True),
nn.Dropout3d(p=0.2),
)
self.block4 = nn.Sequential(
nn.Conv3d(256,
256,
kernel_size=(3, 3, 3),
stride=1,
dilation=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(256),
nn.ReLU(inplace=True),
nn.Conv3d(256,
256,
kernel_size=(3, 3, 3),
stride=1,
dilation=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(256),
nn.ReLU(inplace=True),
nn.Conv3d(256,
512,
kernel_size=(3, 3, 3),
stride=(1, 2, 2),
dilation=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(512),
nn.ReLU(inplace=True),
nn.Dropout3d(p=0.2),
)
self.block5 = nn.Sequential(
nn.Conv3d(512,
512,
kernel_size=(3, 3, 3),
stride=1,
dilation=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(512),
nn.ReLU(inplace=True),
nn.Conv3d(512,
512,
kernel_size=(3, 3, 3),
stride=(1, 2, 2),
dilation=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(512),
nn.ReLU(inplace=True),
)
def __init__(self, in_channels, n_classes, base_n_filter=8):
super(Modified3DUNet, self).__init__()
self.in_channels = in_channels
self.n_classes = n_classes
self.base_n_filter = base_n_filter
self.lrelu = nn.LeakyReLU()
self.dropout3d = nn.Dropout3d(p=0.6)
self.upsacle = nn.Upsample(scale_factor=2, mode='nearest')
self.softmax = nn.Softmax(dim=1)
# Level 1 context pathway
self.conv3d_c1_1 = nn.Conv3d(self.in_channels, self.base_n_filter, kernel_size=3, stride=1, padding=1,
bias=False)
self.conv3d_c1_2 = nn.Conv3d(self.base_n_filter, self.base_n_filter, kernel_size=3, stride=1, padding=1,
bias=False)
self.lrelu_conv_c1 = self.lrelu_conv(self.base_n_filter, self.base_n_filter)
self.inorm3d_c1 = nn.InstanceNorm3d(self.base_n_filter)
# Level 2 context pathway
self.conv3d_c2 = nn.Conv3d(self.base_n_filter, self.base_n_filter * 2, kernel_size=3, stride=2, padding=1,
bias=False)
self.norm_lrelu_conv_c2 = self.norm_lrelu_conv(self.base_n_filter * 2, self.base_n_filter * 2)
self.inorm3d_c2 = nn.InstanceNorm3d(self.base_n_filter * 2)
# Level 3 context pathway
self.conv3d_c3 = nn.Conv3d(self.base_n_filter * 2, self.base_n_filter * 4, kernel_size=3, stride=2, padding=1,
bias=False)
self.norm_lrelu_conv_c3 = self.norm_lrelu_conv(self.base_n_filter * 4, self.base_n_filter * 4)
self.inorm3d_c3 = nn.InstanceNorm3d(self.base_n_filter * 4)
# Level 4 context pathway
self.conv3d_c4 = nn.Conv3d(self.base_n_filter * 4, self.base_n_filter * 8, kernel_size=3, stride=2, padding=1,
bias=False)
self.norm_lrelu_conv_c4 = self.norm_lrelu_conv(self.base_n_filter * 8, self.base_n_filter * 8)
self.inorm3d_c4 = nn.InstanceNorm3d(self.base_n_filter * 8)
# Level 5 context pathway, level 0 localization pathway
self.conv3d_c5 = nn.Conv3d(self.base_n_filter * 8, self.base_n_filter * 16, kernel_size=3, stride=2, padding=1,
bias=False)
self.norm_lrelu_conv_c5 = self.norm_lrelu_conv(self.base_n_filter * 16, self.base_n_filter * 16)
self.norm_lrelu_upscale_conv_norm_lrelu_l0 = self.norm_lrelu_upscale_conv_norm_lrelu(self.base_n_filter * 16,
self.base_n_filter * 8)
self.conv3d_l0 = nn.Conv3d(self.base_n_filter * 8, self.base_n_filter * 8, kernel_size=1, stride=1, padding=0,
bias=False)
self.inorm3d_l0 = nn.InstanceNorm3d(self.base_n_filter * 8)
# Level 1 localization pathway
self.conv_norm_lrelu_l1 = self.conv_norm_lrelu(self.base_n_filter * 16, self.base_n_filter * 16)
self.conv3d_l1 = nn.Conv3d(self.base_n_filter * 16, self.base_n_filter * 8, kernel_size=1, stride=1, padding=0,
bias=False)
self.norm_lrelu_upscale_conv_norm_lrelu_l1 = self.norm_lrelu_upscale_conv_norm_lrelu(self.base_n_filter * 8,
self.base_n_filter * 4)
# Level 2 localization pathway
self.conv_norm_lrelu_l2 = self.conv_norm_lrelu(self.base_n_filter * 8, self.base_n_filter * 8)
self.conv3d_l2 = nn.Conv3d(self.base_n_filter * 8, self.base_n_filter * 4, kernel_size=1, stride=1, padding=0,
bias=False)
self.norm_lrelu_upscale_conv_norm_lrelu_l2 = self.norm_lrelu_upscale_conv_norm_lrelu(self.base_n_filter * 4,
self.base_n_filter * 2)
# Level 3 localization pathway
self.conv_norm_lrelu_l3 = self.conv_norm_lrelu(self.base_n_filter * 4, self.base_n_filter * 4)
self.conv3d_l3 = nn.Conv3d(self.base_n_filter * 4, self.base_n_filter * 2, kernel_size=1, stride=1, padding=0,
bias=False)
self.norm_lrelu_upscale_conv_norm_lrelu_l3 = self.norm_lrelu_upscale_conv_norm_lrelu(self.base_n_filter * 2,
self.base_n_filter)
# Level 4 localization pathway
self.conv_norm_lrelu_l4 = self.conv_norm_lrelu(self.base_n_filter * 2, self.base_n_filter * 2)
self.conv3d_l4 = nn.Conv3d(self.base_n_filter * 2, self.n_classes, kernel_size=1, stride=1, padding=0,
bias=False)
self.ds2_1x1_conv3d = nn.Conv3d(self.base_n_filter * 8, self.n_classes, kernel_size=1, stride=1, padding=0,
bias=False)
self.ds3_1x1_conv3d = nn.Conv3d(self.base_n_filter * 4, self.n_classes, kernel_size=1, stride=1, padding=0,
bias=False)
def __init__(self):
super(Net, self).__init__()
# The first few layers consumes the most memory, so use simple convolution to save memory.
# Call these layers preBlock, i.e., before the residual blocks of later layers.
self.preBlock = nn.Sequential(
nn.Conv3d(1, 24, kernel_size=3, padding=1), nn.BatchNorm3d(24),
nn.ReLU(inplace=True), nn.Conv3d(24, 24, kernel_size=3, padding=1),
nn.BatchNorm3d(24), nn.ReLU(inplace=True))
# 3 poolings, each pooling downsamples the feature map by a factor 2.
# 3 groups of blocks. The first block of each group has one pooling.
num_blocks_forw = [2, 2, 3, 3] # 前向传播的完整网络
num_blocks_back = [3, 3] # 反向传播需要更新参数的blocks（相当于keras的trainable）
self.featureNum_forw = [24, 32, 64, 64, 64] # ？
self.featureNum_back = [128, 64, 64] # ？
for i in range(len(num_blocks_forw)):
blocks = []
for j in range(num_blocks_forw[i]):
if j == 0:
blocks.append(
PostRes(self.featureNum_forw[i],
self.featureNum_forw[i + 1]))
else:
blocks.append(
PostRes(self.featureNum_forw[i + 1],
self.featureNum_forw[i + 1]))
setattr(self, 'forw' + str(i + 1), nn.Sequential(*blocks))
for i in range(len(num_blocks_back)):
blocks = []
for j in range(num_blocks_back[i]):
if j == 0:
if i == 0:
addition = 3
else:
addition = 0
blocks.append(
PostRes(
self.featureNum_back[i + 1] +
self.featureNum_forw[i + 2] + addition,
self.featureNum_back[i]))
else:
blocks.append(
PostRes(self.featureNum_back[i],
self.featureNum_back[i]))
setattr(self, 'back' + str(i + 2), nn.Sequential(*blocks))
self.maxpool1 = nn.MaxPool3d(kernel_size=2,
stride=2,
return_indices=True)
self.maxpool2 = nn.MaxPool3d(kernel_size=2,
stride=2,
return_indices=True)
self.maxpool3 = nn.MaxPool3d(kernel_size=2,
stride=2,
return_indices=True)
self.maxpool4 = nn.MaxPool3d(kernel_size=2,
stride=2,
return_indices=True)
self.unmaxpool1 = nn.MaxUnpool3d(kernel_size=2, stride=2)
self.unmaxpool2 = nn.MaxUnpool3d(kernel_size=2, stride=2)
self.path1 = nn.Sequential(
nn.ConvTranspose3d(64, 64, kernel_size=2, stride=2),
nn.BatchNorm3d(64), nn.ReLU(inplace=True))
self.path2 = nn.Sequential(
nn.ConvTranspose3d(64, 64, kernel_size=2, stride=2),
nn.BatchNorm3d(64), nn.ReLU(inplace=True))
self.drop = nn.Dropout3d(p=0.5, inplace=False)
self.output = nn.Sequential(
nn.Conv3d(self.featureNum_back[0], 64, kernel_size=1),
nn.ReLU(),
# nn.Dropout3d(p = 0.3),
nn.Conv3d(64, 5 * len(config['anchors']), kernel_size=1))
def __init__(
self,
maxpool,
in_channels,
out_channels,
kernel_sizes,
strides,
dilatations,
padding,
batchnorm,
n_classes,
is_bayesian,
max_fvc,
n_kernels,
random_node='output',
activation=torch.nn.ReLU,
n_res=3,
gated=True,
has_dense=True,
resblocks=False,
):
super().__init__()
self.max_fvc = max_fvc
if torch.cuda.is_available():
device = 'cuda'
else:
device = 'cpu'
self.maxpool = nn.MaxPool3d(maxpool, return_indices=False)
self.is_bayesian = is_bayesian
if is_bayesian:
self.GaussianSample = GaussianUnitSample(37, 37)
self.GaussianBlock = GaussianBlock(1, 1, n_kernels=n_kernels)
self.device = device
self.conv_layers = nn.ModuleList()
self.deconv_layers = nn.ModuleList()
self.bns = nn.ModuleList()
self.dropout = nn.ModuleList()
self.resconv = nn.ModuleList()
self.activation = activation()
self.n_res = n_res
self.resblocks = resblocks
self.has_dense = has_dense
self.batchnorm = batchnorm
self.a_dim = None
for i, (ins, outs, ksize, stride, dilats, pad) in enumerate(
zip(in_channels, out_channels, kernel_sizes, strides,
dilatations, padding)):
if not gated:
self.conv_layers += [
torch.nn.Conv3d(
in_channels=ins,
out_channels=outs,
kernel_size=ksize,
stride=stride,
padding=pad,
dilation=dilats,
)
]
else:
self.conv_layers += [
GatedConv3d(input_channels=ins,
output_channels=outs,
kernel_size=ksize,
stride=stride,
padding=pad,
dilation=dilats,
activation=nn.Tanh())
]
if resblocks and i != 0:
for _ in range(n_res):
self.resconv += [ResBlock3D(ins, outs, activation, device)]
self.bns += [nn.BatchNorm3d(num_features=outs)]
self.dropout3d = nn.Dropout3d(0.5)
self.dense1 = torch.nn.Linear(in_features=out_channels[-1],
out_features=32)
self.dense1_bn = nn.BatchNorm1d(num_features=32)
self.dense2 = torch.nn.Linear(
in_features=32 + 5,
out_features=n_classes) # 5 parameters added here
self.dense2_bn = nn.BatchNorm1d(num_features=n_classes)
self.dropout = nn.Dropout(0.5)
self.log_softmax = torch.nn.functional.log_softmax
def __init__(self,
in_channel=1,
out_channel=3,
filters=[8, 12, 16, 20, 24],
act='sigmoid',
aux=False):
#def __init__(self, in_channel=1, out_channel=3, filters=[28, 36, 48, 64, 80]):
super().__init__()
# encoding path
self.layer1_E = nn.Sequential(
conv3d_bn_elu(in_planes=in_channel,
out_planes=filters[0],
kernel_size=(1, 5, 5),
stride=1,
padding=(0, 2, 2)),
conv3d_bn_elu(in_planes=filters[0],
out_planes=filters[0],
kernel_size=(1, 3, 3),
stride=1,
padding=(0, 1, 1)),
residual_block_2d(filters[0], filters[0], projection=False))
self.layer2_E = nn.Sequential(
conv3d_bn_elu(in_planes=filters[0],
out_planes=filters[1],
kernel_size=(1, 3, 3),
stride=1,
padding=(0, 1, 1)),
residual_block_3d(filters[1], filters[1], projection=False))
self.layer3_E = nn.Sequential(
conv3d_bn_elu(in_planes=filters[1],
out_planes=filters[2],
kernel_size=(1, 3, 3),
stride=1,
padding=(0, 1, 1)),
residual_block_3d(filters[2], filters[2], projection=False))
self.layer4_E = nn.Sequential(
conv3d_bn_elu(in_planes=filters[2],
out_planes=filters[3],
kernel_size=(1, 3, 3),
stride=1,
padding=(0, 1, 1)),
residual_block_3d(filters[3], filters[3], projection=False))
# center block
self.center = nn.Sequential(
conv3d_bn_elu(in_planes=filters[3],
out_planes=filters[4],
kernel_size=(1, 3, 3),
stride=1,
padding=(0, 1, 1)),
residual_block_3d(filters[4], filters[4], projection=True))
# decoding path
self.layer1_D = nn.Sequential(
conv3d_bn_elu(in_planes=filters[0],
out_planes=filters[0],
kernel_size=(1, 3, 3),
stride=1,
padding=(0, 1, 1)),
residual_block_2d(filters[0], filters[0], projection=False),
conv3d_bn_non(in_planes=filters[0],
out_planes=out_channel,
kernel_size=(1, 5, 5),
stride=1,
padding=(0, 2, 2)))
self.layer2_D = nn.Sequential(
conv3d_bn_elu(in_planes=filters[1],
out_planes=filters[1],
kernel_size=(1, 3, 3),
stride=1,
padding=(0, 1, 1)),
residual_block_3d(filters[1], filters[1], projection=False))
self.layer3_D = nn.Sequential(
conv3d_bn_elu(in_planes=filters[2],
out_planes=filters[2],
kernel_size=(1, 3, 3),
stride=1,
padding=(0, 1, 1)),
residual_block_3d(filters[2], filters[2], projection=False))
self.layer4_D = nn.Sequential(
conv3d_bn_elu(in_planes=filters[3],
out_planes=filters[3],
kernel_size=(1, 3, 3),
stride=1,
padding=(0, 1, 1)),
residual_block_3d(filters[3], filters[3], projection=False))
# pooling & upsample
self.down = nn.MaxPool3d(kernel_size=(1, 2, 2), stride=(1, 2, 2))
self.up = nn.Upsample(scale_factor=(1, 2, 2),
mode='trilinear',
align_corners=False)
# conv + upsample
self.conv1 = conv3d_bn_elu(filters[1],
filters[0],
kernel_size=(1, 1, 1),
padding=(0, 0, 0))
self.conv2 = conv3d_bn_elu(filters[2],
filters[1],
kernel_size=(1, 1, 1),
padding=(0, 0, 0))
self.conv3 = conv3d_bn_elu(filters[3],
filters[2],
kernel_size=(1, 1, 1),
padding=(0, 0, 0))
self.conv4 = conv3d_bn_elu(filters[4],
filters[3],
kernel_size=(1, 1, 1),
padding=(0, 0, 0))
if act == 'tanh':
self.act = nn.Tanh()
else:
self.act = nn.Sigmoid()
self.aux = aux
if self.aux == 1:
self.dsn = nn.Sequential(
conv3d_bn_elu(filters[2],
filters[2] // 2,
kernel_size=(1, 3, 3),
stride=1,
padding=(0, 1, 1)), nn.Dropout3d(0.1),
nn.Conv3d(filters[2] // 2,
out_channel,
kernel_size=(1, 1, 1),
stride=1,
padding=(0, 0, 0),
bias=True))
else:
self.dsn = None
#initialization
ortho_init(self)
def __init__(self,
n_inp=1,
feats=[32, 64, 64, 128, 128, 256, 256],
blocks=[2, 2, 2, 2, 2, 2, 2],
n_pred_p=[1, 1, 1],
n_pred_d=[4, 4, 4],
L_output=[0, 1, 2],
abn=1,
dropout_ratio=0.0):
super().__init__()
self.relu = nn.ReLU(inplace=True)
self.pixel_shuffle = nn.PixelShuffle(2)
self.upsample = nn.Upsample(scale_factor=2, mode='nearest')
self.n_inp = n_inp
if abn == 0:
abnblock = ABN
elif abn == 1:
abnblock = InPlaceABN
else:
abnblock = InPlaceABNSync
Conv3d_ABN.define_abn(abnblock)
Bottleneck_ABN.define_abn(abnblock)
self.strides = np.array([[1, 1, 1], [2, 2, 2], [4, 4, 4], [8, 8, 8],
[16, 16, 16]])
for i in range(6):
if i == 0:
setattr(self, 'l' + str(i),
build_layer(n_inp, feats[i], blocks[i]))
setattr(self, 'r' + str(i),
build_layer(feats[i], feats[i], blocks[i], skip=True))
else:
setattr(
self, 'l' + str(i),
build_layer(feats[i - 1], feats[i], blocks[i], stride=2))
setattr(
self, 'rnnlz' + str(i),
nn.LSTM(feats[i], feats[i] // 2, 1, bidirectional=True))
setattr(self, 'rnnlz_bn' + str(i), abnblock(feats[i]))
setattr(
self, 'rnnlx' + str(i),
nn.LSTM(feats[i], feats[i] // 2, 1, bidirectional=True))
setattr(self, 'rnnlx_bn' + str(i), abnblock(feats[i]))
setattr(
self, 'rnnly' + str(i),
nn.LSTM(feats[i], feats[i] // 2, 1, bidirectional=True))
setattr(self, 'rnnly_bn' + str(i), abnblock(feats[i]))
setattr(self, 'r' + str(i),
build_layer(feats[i], feats[i], blocks[i], skip=True))
setattr(self, 'rt' + str(i),
build_layer(feats[i], feats[i - 1], blocks[i]))
self.r0 = Conv3d_ABN(feats[0], feats[0], kernel_size=3, padding=1)
self.out_layer = nn.Conv3d(feats[0], 1, kernel_size=1, stride=1)
self.drop_out = nn.Dropout3d(p=dropout_ratio)
for c_out1, c_out2, out in zip(n_pred_p, n_pred_d, L_output):
#outp = build_layer(feats[out], feats[out] // 4, 0, skip=False)
outp = nn.Conv3d(feats[out], feats[out], kernel_size=3, padding=1)
setattr(self, 'before_out' + str(out) + 'p', outp)
#outd = build_layer(feats[out], feats[out] // 4, 0, skip=False)
outd = nn.Conv3d(feats[out], feats[out], kernel_size=3, padding=1)
setattr(self, 'before_out' + str(out) + 'd', outd)
for c_out1, c_out2, out in zip(n_pred_p, n_pred_d, L_output):
setattr(self, 'out' + str(out) + 'p',
nn.Conv3d(feats[out], c_out1, kernel_size=3, padding=1))
outd = nn.Conv3d(feats[out], c_out2, kernel_size=3, padding=1)
#outd.weight.data.fill_(0.0)
outd.bias.data.fill_(0.0)
setattr(self, 'out' + str(out) + 'd', outd)
self.L_output = L_output
self.strides = self.strides[np.array(L_output)]
def __init__(
self,
n_classes=2,
in_channels=3,
norm_type='GN_8', # default group norm with 8 groups
):
super(RSANet, self).__init__()
filters = [32, 64, 128, 256, 512]
self.init_conv = nn.Conv3d(in_channels,
filters[0],
3,
1,
1,
bias=False)
self.dropout = nn.Dropout3d(0.2)
self.conv1 = unetConvBlock(filters[0],
filters[0],
norm_type,
num_convs=2)
self.down1 = unetDownSample(filters[0],
down_type='conv',
norm_type=norm_type)
self.conv2 = unetConvBlock(filters[0],
filters[1],
norm_type,
num_convs=2)
self.down2 = unetDownSample(filters[1],
down_type='conv',
norm_type=norm_type)
self.conv3 = unetConvBlock(filters[1],
filters[2],
norm_type,
num_convs=2)
self.down3 = unetDownSample(filters[2],
down_type='conv',
norm_type=norm_type)
self.conv4 = unetConvBlock(filters[2],
filters[3],
norm_type,
num_convs=2)
self.down4 = unetDownSample(filters[3],
down_type='conv',
norm_type=norm_type)
self.center = unetConvBlock(filters[3],
filters[4],
norm_type,
num_convs=2)
self.up_concat4 = unetUpPadCatConv(filters[3], filters[4], False,
norm_type)
self.up_concat3 = unetUpPadCatConv(filters[2], filters[3], False,
norm_type)
self.up_concat2 = unetUpPadCatConv(filters[1], filters[2], False,
norm_type)
self.up_concat1 = unetUpPadCatConv(filters[0], filters[1], False,
norm_type)
self.final = nn.Sequential(
nn.Conv3d(filters[0], filters[0], 3, 1, 1, bias=False),
nn.Conv3d(filters[0], n_classes, 1, bias=False),
)
# --------- Recurrent slice-wise attention (RSA) module --------- #
self.rsa_block = rsaBlock(filters[4])
def __init__(self,
input_channels,
feature_channels,
n_stack,
n_joint,
downsample=1,
configs=None):
super(V2V_HG, self).__init__()
self.input_channels = input_channels
self.n_stack = n_stack
self.n_joint = n_joint
self.configs = configs
# fix feature_channels to n_joitn(=22) like in SkelVolNet
# feature_channels = n_joint
self.feature_channels = feature_channels
if configs['HG_type'] == 'double':
HG = HG_double
elif configs['HG_type'] == 'double2':
HG = HG_double
elif configs['HG_type'] == 'double2_attention':
HG = HG_double2_attention
if downsample > 1:
self.front_layers = nn.Sequential(
Basic3DBlock(input_channels, feature_channels, 5),
Res3DBlock(feature_channels, feature_channels),
Pool3DBlock(int(downsample), feature_channels))
else:
self.front_layers = nn.Sequential(
Basic3DBlock(input_channels, feature_channels, 5),
Res3DBlock(feature_channels, feature_channels))
self.hg_1 = HG(input_channels=feature_channels,
output_channels=feature_channels,
N=88 // downsample)
self.joint_output_1 = nn.Sequential(
Res3DBlock(feature_channels, feature_channels // 2),
Basic3DBlock(feature_channels // 2, feature_channels // 2, 1),
nn.Dropout3d(p=0.2),
nn.Conv3d(feature_channels // 2,
n_joint,
kernel_size=1,
stride=1,
padding=0))
if n_stack > 1:
self.hg_list = nn.ModuleList([
HG(input_channels=feature_channels + n_joint,
output_channels=feature_channels,
N=88 // downsample) for i in range(1, n_stack)
])
self.joint_output_list = nn.ModuleList([
nn.Sequential(
Res3DBlock(feature_channels, feature_channels // 2),
Basic3DBlock(feature_channels // 2, feature_channels // 2,
1), nn.Dropout3d(p=0.2),
nn.Conv3d(feature_channels // 2,
n_joint,
kernel_size=1,
stride=1,
padding=0)) for i in range(1, n_stack)
])
self._initialize_weights()
