def __init__(self, in_ch, out_ch):
super(up_2_2_1, self).__init__()
self.up = nn.ConvTranspose3d(in_ch, out_ch, kernel_size=(3, 3, 3), stride=(1, 2, 2), padding=(1, 1, 1), output_padding=(0,1,1))
def __init__(self, in_channels=3, n_sequence=3, out_channels=3, n_resblock=3, n_feat=32, device='cuda'):
super(HybridC3D, self).__init__()
print("Creating HybridC3D Net")
self.n_sequence = n_sequence
self.device = device
assert n_sequence == 3, "Only support args.n_sequence=3; but get args.n_sequence={}".format(n_sequence)
InBlock = []
# b, 3, 3, 256, 256
InBlock.extend([nn.Sequential(
nn.Conv3d(in_channels, n_feat, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1)),
nn.ReLU(inplace=True)
)])
# b, 32, 3, 256, 256
InBlock.extend([blocks.ResBlock3D(n_feat, n_feat, kernel_size=(3, 3, 3), stride=(1, 1, 1))
for _ in range(n_resblock)])
# b, 32, 3, 256, 256
# encoder1
Encoder_first = [nn.Sequential(
nn.Conv3d(n_feat, n_feat * 2, kernel_size=(3, 3, 3), stride=(1, 2, 2), padding=(1, 1, 1)),
nn.ReLU(inplace=True)
)]
# b, 64, 3, 128, 128
Encoder_first.extend([blocks.ResBlock3D(n_feat * 2, n_feat * 2, kernel_size=(3, 3, 3), stride=(1, 1, 1))
for _ in range(n_resblock)])
# b, 64, 3, 128, 128
# encoder2
Encoder_second = [nn.Sequential(
nn.Conv3d(n_feat * 2, n_feat * 4, kernel_size=(3, 3, 3), stride=(1, 2, 2), padding=(0, 1, 1)),
nn.ReLU(inplace=True)
)]
# b, 128, 1, 64, 64
Encoder_second.extend([blocks.ResBlock3D(n_feat * 4, n_feat * 4, kernel_size=(1, 3, 3),
stride=(1, 1, 1), padding=(0, 1, 1)) for _ in range(n_resblock)])
# decoder2
Decoder_second = [nn.Sequential(
nn.ConvTranspose3d(n_feat * 4, n_feat * 2, kernel_size=(1, 3, 3), stride=(1, 2, 2),
padding=(0, 1, 1), output_padding=(0, 1, 1)),
nn.ReLU(inplace=True)
)]
Decoder_second.extend([blocks.ResBlock3D(n_feat * 2, n_feat * 2, kernel_size=(1, 3, 3),
stride=(1, 1, 1), padding=(0, 1, 1)) for _ in range(n_resblock)])
# decoder1
Decoder_first = [nn.Sequential(
nn.ConvTranspose3d(n_feat * 2, n_feat, kernel_size=(1, 3, 3), stride=(1, 2, 2),
padding=(0, 1, 1), output_padding=(0, 1, 1)),
nn.ReLU(inplace=True)
)]
Decoder_first.extend([blocks.ResBlock3D(n_feat, n_feat, kernel_size=(1, 3, 3),
stride=(1, 1, 1), padding=(0, 1, 1)) for _ in range(n_resblock)])
OutBlock = nn.Conv3d(n_feat, out_channels, kernel_size=(1, 3, 3), stride=(1, 1, 1), padding=(0, 1, 1))
self.inBlock = nn.Sequential(*InBlock)
self.encoder_first = nn.Sequential(*Encoder_first)
self.encoder_second = nn.Sequential(*Encoder_second)
self.decoder_second = nn.Sequential(*Decoder_second)
self.decoder_first = nn.Sequential(*Decoder_first)
self.outBlock = OutBlock
def _conv3d_tro(self, in_channels, out_channels):
return nn.ConvTranspose3d(in_channels=in_channels,
out_channels=out_channels,
kernel_size=(1, 1, 1),
stride=(2, 1, 1),
padding=(0, 0, 0))
def __init__(self, in_ch, out_ch, attention=False):
super(MyselfUnet3d, self).__init__()
self.conv0 = nn.Sequential(nn.Conv3d(in_ch, 32, 3, padding=1),
nn.ReLU(inplace=True),
nn.Conv3d(32, 64, 3, padding=1),
nn.ReLU(inplace=True))
self.attention = attention
self.pool1 = nn.MaxPool3d((1, 2, 2),
(1, 2, 2)) # (kernel_size, stride)
self.conv1 = downDouble3dConv(64, 128)
self.pool2 = nn.MaxPool3d((1, 2, 2), (1, 2, 2))
self.conv2 = downDouble3dConv(128, 256)
self.pool3 = nn.MaxPool3d((1, 2, 2), (1, 2, 2))
self.spp_convc1 = nn.Conv3d(128,
16,
kernel_size=3,
dilation=1,
padding=1)
# self.spp_convp1 = nn.Conv3d(64, 16, 1)
self.spp_convc2 = nn.Conv3d(256,
16,
kernel_size=3,
dilation=3,
padding=3)
self.spp_up_c1 = nn.ConvTranspose3d(16,
16,
kernel_size=(1, 2, 2),
stride=(1, 2, 2))
self.spp_up_c2 = nn.ConvTranspose3d(16,
16,
kernel_size=(1, 4, 4),
stride=(1, 4, 4))
# self.spp_up_p3 = nn.ConvTranspose3d(16, 16, kernel_size = (1, 8, 8), stride=(1, 8, 8))
self.weight_cat_all3 = SElayer(96)
self.refine3 = nn.Conv3d(96, 96, 3, padding=1)
self.spp_convc22 = nn.Conv3d(256,
16,
kernel_size=3,
dilation=1,
padding=1)
# self.spp_convc32 = nn.Conv3d(256, 16, kernel_size = 3, dilation = 3, padding=3)
self.spp_up_c22 = nn.ConvTranspose3d(16,
16,
kernel_size=(1, 2, 2),
stride=(1, 2, 2))
# self.spp_up_p32 = nn.ConvTranspose3d(16, 16, kernel_size = (1, 4, 4), stride=(1, 4, 4))
self.weight_cat_all2 = SElayer(144)
self.refine2 = nn.Conv3d(144, 144, 3, padding=1)
self.spp_convp33 = nn.Conv3d(256,
16,
kernel_size=3,
dilation=1,
padding=1)
self.spp_up_p33 = nn.ConvTranspose3d(16,
16,
kernel_size=(1, 2, 2),
stride=(1, 2, 2))
self.bridge = downDouble3dConv(256, 512)
# 输出level3的预测结果
self.output_l3 = nn.Conv3d(512, 1, 3, padding=1)
self.up1 = nn.ConvTranspose3d(512, 512, (1, 2, 2), stride=(1, 2, 2))
self.conv4 = upDouble3dConv(768, 256)
# self.output_l2 = nn.Conv3d(256)
self.up2 = nn.ConvTranspose3d(256, 256, (1, 2, 2), stride=(1, 2, 2))
self.conv5 = upDouble3dConv(400, 128)
self.up3 = nn.ConvTranspose3d(128, 128, (1, 2, 2),
stride=(1, 2, 2)) ##
self.conv6 = upDouble3dConv(224, 64)
self.conv7 = nn.Conv3d(64, 1, 3, padding=1)
self.BN3d = nn.BatchNorm3d(out_ch)
def lua_recursive_model(module,seq):
for m in module.modules:
name = type(m).__name__
real = m
if name == 'TorchObject':
name = m._typename.replace('cudnn.','')
m = m._obj
if name == 'SpatialConvolution' or name == 'nn.SpatialConvolutionMM':
if not hasattr(m,'groups') or m.groups is None: m.groups=1
n = nn.Conv2d(m.nInputPlane,m.nOutputPlane,(m.kW,m.kH),(m.dW,m.dH),(m.padW,m.padH),1,m.groups,bias=(m.bias is not None))
copy_param(m,n)
add_submodule(seq,n)
elif name == 'SpatialBatchNormalization':
n = nn.BatchNorm2d(m.running_mean.size(0), m.eps, m.momentum, m.affine)
copy_param(m,n)
add_submodule(seq,n)
elif name == 'VolumetricBatchNormalization':
n = nn.BatchNorm3d(m.running_mean.size(0), m.eps, m.momentum, m.affine)
copy_param(m, n)
add_submodule(seq, n)
elif name == 'ReLU':
n = nn.ReLU()
add_submodule(seq,n)
elif name == 'Sigmoid':
n = nn.Sigmoid()
add_submodule(seq,n)
elif name == 'SpatialMaxPooling':
n = nn.MaxPool2d((m.kW,m.kH),(m.dW,m.dH),(m.padW,m.padH),ceil_mode=m.ceil_mode)
add_submodule(seq,n)
elif name == 'SpatialAveragePooling':
n = nn.AvgPool2d((m.kW,m.kH),(m.dW,m.dH),(m.padW,m.padH),ceil_mode=m.ceil_mode)
add_submodule(seq,n)
elif name == 'SpatialUpSamplingNearest':
n = nn.UpsamplingNearest2d(scale_factor=m.scale_factor)
add_submodule(seq,n)
elif name == 'View':
n = Lambda(lambda x: x.view(x.size(0),-1))
add_submodule(seq,n)
elif name == 'Reshape':
n = Lambda(lambda x: x.view(x.size(0),-1))
add_submodule(seq,n)
elif name == 'Linear':
# Linear in pytorch only accept 2D input
n1 = Lambda(lambda x: x.view(1,-1) if 1==len(x.size()) else x )
n2 = nn.Linear(m.weight.size(1),m.weight.size(0),bias=(m.bias is not None))
copy_param(m,n2)
n = nn.Sequential(n1,n2)
add_submodule(seq,n)
elif name == 'Dropout':
m.inplace = False
n = nn.Dropout(m.p)
add_submodule(seq,n)
elif name == 'SoftMax':
n = nn.Softmax()
add_submodule(seq,n)
elif name == 'Identity':
n = Lambda(lambda x: x) # do nothing
add_submodule(seq,n)
elif name == 'SpatialFullConvolution':
n = nn.ConvTranspose2d(m.nInputPlane,m.nOutputPlane,(m.kW,m.kH),(m.dW,m.dH),(m.padW,m.padH),(m.adjW,m.adjH))
copy_param(m,n)
add_submodule(seq,n)
elif name == 'VolumetricFullConvolution':
n = nn.ConvTranspose3d(m.nInputPlane,m.nOutputPlane,(m.kT,m.kW,m.kH),(m.dT,m.dW,m.dH),(m.padT,m.padW,m.padH),(m.adjT,m.adjW,m.adjH),m.groups)
copy_param(m,n)
add_submodule(seq, n)
elif name == 'SpatialReplicationPadding':
n = nn.ReplicationPad2d((m.pad_l,m.pad_r,m.pad_t,m.pad_b))
add_submodule(seq,n)
elif name == 'SpatialReflectionPadding':
n = nn.ReflectionPad2d((m.pad_l,m.pad_r,m.pad_t,m.pad_b))
add_submodule(seq,n)
elif name == 'Copy':
n = Lambda(lambda x: x) # do nothing
add_submodule(seq,n)
elif name == 'Narrow':
n = Lambda(lambda x,a=(m.dimension,m.index,m.length): x.narrow(*a))
add_submodule(seq,n)
elif name == 'SpatialCrossMapLRN':
lrn = lnn.SpatialCrossMapLRN(m.size,m.alpha,m.beta,m.k)
n = Lambda(lambda x,lrn=lrn: Variable(lrn.forward(x.data)))
add_submodule(seq,n)
elif name == 'Sequential':
n = nn.Sequential()
lua_recursive_model(m,n)
add_submodule(seq,n)
elif name == 'ConcatTable': # output is list
n = LambdaMap(lambda x: x)
lua_recursive_model(m,n)
add_submodule(seq,n)
elif name == 'CAddTable': # input is list
n = LambdaReduce(lambda x,y: x+y)
add_submodule(seq,n)
elif name == 'Concat':
dim = m.dimension
n = LambdaReduce(lambda x,y,dim=dim: torch.cat((x,y),dim))
lua_recursive_model(m,n)
add_submodule(seq,n)
elif name == 'TorchObject':
print('Not Implement',name,real._typename)
else:
print('Not Implement',name)
def __init__(
self,
layer_1_cnn_filters,
layer_2_cnn_filters,
layer_3_cnn_filters,
layer_4_cnn_filters,
layer_1_kernel_size,
layer_2_kernel_size,
layer_3_kernel_size,
layer_4_kernel_size,
force_hidden_layer_size,
middle_hidden_layer_size,
recurrent_layer_size,
device,
):
"""
force_add determines whether the force is added or concatonated.
"""
super(RecurrentModel, self).__init__()
self.layer_1_cnn_filters = layer_1_cnn_filters
self.layer_2_cnn_filters = layer_2_cnn_filters
self.layer_3_cnn_filters = layer_3_cnn_filters
self.layer_4_cnn_filters = layer_4_cnn_filters
self.layer_1_kernel_size = layer_1_kernel_size
self.layer_2_kernel_size = layer_2_kernel_size
self.layer_3_kernel_size = layer_3_kernel_size
self.layer_4_kernel_size = layer_4_kernel_size
self.middle_hidden_layer_size = middle_hidden_layer_size
self.recurrent_layer_size = recurrent_layer_size
self.device = device
self.init_recurrent_state = (
torch.nn.Parameter(torch.rand(LSTM_DEPTH, 1,
middle_hidden_layer_size),
requires_grad=True).to(self.device),
torch.nn.Parameter(torch.rand(LSTM_DEPTH, 1,
middle_hidden_layer_size),
requires_grad=True).to(self.device),
)
LAYERS = 4
self.middle_layer_image_width = int(GRID_SIZE / (2**(LAYERS - 1)))
middle_layer_size = int(layer_4_cnn_filters * 1 *
(self.middle_layer_image_width**2))
self.middle_layer_size = middle_layer_size
self.layer_tcnn_0 = nn.Sequential(
nn.ConvTranspose3d(
layer_4_cnn_filters,
layer_3_cnn_filters,
kernel_size=layer_4_kernel_size,
stride=[STRIDE, STRIDE, STRIDE],
padding=int(layer_4_kernel_size / 3),
output_padding=[1, 1, 1],
),
# nn.BatchNorm2d(4),
)
self.layer_tcnn_1 = nn.Sequential(
nn.ConvTranspose3d(
layer_3_cnn_filters,
layer_2_cnn_filters,
kernel_size=layer_3_kernel_size,
stride=[STRIDE, STRIDE, STRIDE],
padding=int(layer_3_kernel_size / 2),
output_padding=[1, 1, 1],
),
# nn.BatchNorm2d(4),
)
self.layer_tcnn_2 = nn.Sequential(
nn.ConvTranspose3d(
layer_2_cnn_filters,
layer_1_cnn_filters,
kernel_size=layer_2_kernel_size,
stride=[STRIDE, STRIDE, 1],
padding=int(layer_2_kernel_size / 2),
output_padding=[1, 1, 0],
),
# nn.BatchNorm2d(2),
)
self.layer_tcnn_3 = nn.Sequential(
nn.ConvTranspose3d(
layer_1_cnn_filters,
IMAGE_DEPTH,
kernel_size=layer_1_kernel_size,
stride=[1, 1, 1],
padding=int(layer_1_kernel_size / 2),
output_padding=[0, 0, 0],
),
# nn.BatchNorm2d(4)
)
self.layer_force_recurrent = nn.Sequential(
nn.Linear(4, middle_hidden_layer_size), nn.LeakyReLU())
self.layer_cnn_recurrent = nn.Sequential(
nn.Linear(middle_layer_size, middle_hidden_layer_size),
nn.LeakyReLU())
self.layer_recurrent_out = nn.Sequential(
nn.Linear(middle_hidden_layer_size, middle_layer_size),
nn.LeakyReLU())
self.layer_recurrent = nn.LSTM(middle_hidden_layer_size,
middle_hidden_layer_size, LSTM_DEPTH)
self.layer_sigmoid_out = nn.Sequential(nn.Sigmoid())
self.leaky_relu = nn.LeakyReLU()
def deconv3d_as_up(in_channels, out_channels, kernel_size=2, stride=2):
return nn.Sequential(
nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride),
nn.PReLU())
def __init__(self, in_ch, out_ch):
super(up_2_2_2, self).__init__()
self.up = nn.ConvTranspose3d(in_ch, out_ch, 2, stride=2)
def __init__(self, input_nc, output_nc, depth, ngf=64, use_naive=False):
super(EdsrFGenerator3d, self).__init__()
use_bias = True
downconv_ab = [
nn.ReplicationPad3d(2),
nn.Conv3d(input_nc,
ngf,
kernel_size=5,
stride=1,
padding=0,
bias=use_bias),
nn.InstanceNorm3d(ngf),
nn.ReLU(True),
nn.Conv3d(ngf,
ngf * 2,
kernel_size=3,
stride=2,
padding=1,
bias=use_bias),
nn.InstanceNorm3d(ngf * 2),
nn.ReLU(True)
]
downconv_ba = [
nn.ReplicationPad3d(2),
nn.Conv3d(input_nc,
ngf,
kernel_size=5,
stride=1,
padding=0,
bias=use_bias),
nn.InstanceNorm3d(ngf),
nn.ReLU(True),
nn.Conv3d(ngf,
ngf * 2,
kernel_size=3,
stride=2,
padding=1,
bias=use_bias),
nn.InstanceNorm3d(ngf * 2),
nn.ReLU(True)
]
core = []
for _ in range(depth):
core += [ThickBlock3d(ngf * 2, use_bias, use_naive)]
upconv_ab = [
nn.ConvTranspose3d(ngf * 2,
ngf,
kernel_size=3,
stride=2,
padding=1,
output_padding=1,
bias=use_bias),
nn.InstanceNorm3d(ngf),
nn.ReLU(True),
nn.ReplicationPad3d(2),
nn.Conv3d(ngf, output_nc, kernel_size=5, padding=0),
nn.Tanh()
]
upconv_ba = [
nn.ConvTranspose3d(ngf * 2,
ngf,
kernel_size=3,
stride=2,
padding=1,
output_padding=1,
bias=use_bias),
nn.InstanceNorm3d(ngf),
nn.ReLU(True),
nn.ReplicationPad3d(2),
nn.Conv3d(ngf, output_nc, kernel_size=5, padding=0),
nn.Tanh()
]
self.downconv_ab = nn.Sequential(*downconv_ab)
self.downconv_ba = nn.Sequential(*downconv_ba)
self.core = nn.ModuleList(core)
self.upconv_ab = nn.Sequential(*upconv_ab)
self.upconv_ba = nn.Sequential(*upconv_ba)
def __init__(self, hparams: dict, train_set: Tuple[str, str],
val_set: Tuple[str, str]):
super().__init__()
self.hparams = Namespace(**hparams)
self.vector_length = self.hparams.vector_length
self.dim = self.hparams.dim
self.lr = self.hparams.lr
self.bs = self.hparams.batch_size
self.normal_mean = self.hparams.normal_mean
self.normal_sigma = self.hparams.normal_sigma
self.conv_layers = self.hparams.conv_layers
self.bn = self.hparams.bn
self.elu = self.hparams.elu
self.kld_ratio = self.hparams.kld_ratio
self.weighted_loss = self.hparams.weighted_loss
self.conv_size = validate_conv(self.dim, self.conv_layers)
self.conv_channel = self.conv_layers[-1][1]
if not self.conv_size:
raise RuntimeError("Conv layers not aligned.")
self.encoder = nn.Sequential(
# Example:
# conv_layers=((1, 10, 2, 2, 0), (10, 20, 2, 2, 0))
# nn.Sequential(
# nn.Conv3d(1, 10, 2, 2, 0),
# nn.LeakyRelu(),
# nn.Conv3d(10, 20, 2, 2, 0),
# nn.LeakyRelu()
# )
*(layer for layers in (
[nn.Conv3d(*conv_layer)] +
([nn.BatchNorm3d(conv_layer[1])] if self.bn else []) +
[nn.ELU() if self.elu else nn.ReLU()]
for conv_layer in self.conv_layers) for layer in layers))
self.mu = nn.Sequential(
nn.Linear(
pow(self.conv_size, 3) * self.conv_channel,
self.vector_length), )
self.log_var = nn.Sequential(
nn.Linear(
pow(self.conv_size, 3) * self.conv_channel,
self.vector_length), )
self.fc_decoder = nn.Sequential(
nn.Linear(self.vector_length,
pow(self.conv_size, 3) * self.conv_channel),
nn.ELU() if self.elu else nn.ReLU())
self.decoder = nn.Sequential(
*(layer for layers in ([
nn.ConvTranspose3d(conv_layer[1], conv_layer[0], *
conv_layer[2:])
] + ([nn.BatchNorm3d(conv_layer[0])] if self.bn else []) +
[nn.ELU() if self.elu else nn.ReLU()]
for conv_layer in self.conv_layers[:0:-1])
for layer in layers),
nn.ConvTranspose3d(self.conv_layers[0][1], self.conv_layers[0][0],
*self.conv_layers[0][2:]), nn.Sigmoid())
# Data
# self.train_dataset = NumpyVoxelDataset(train_set[0], self.dim, train_set[1])
# self.val_dataset = NumpyVoxelDataset(val_set[0], self.dim, val_set[1])
# self.val_dataset = PickerDataset(self.train_dataset, ids=[10, 513, 2013, 10403, 20303, 40013])
self.train_dataset = ShapeNetVoxelDataset(
'shapenetvoxel128', 'pix2mesh_splits_val05.json', 128, 'train')
self.val_dataset = ShapeNetVoxelDataset('shapenetvoxel128',
'pix2mesh_splits_val05.json',
128, 'test')
def __init__(self, training):
super().__init__()
self.training = training
self.encoder_stage1 = nn.Sequential(
nn.Conv3d(1, 16, 3, 1, padding=1),
nn.PReLU(16),
nn.Conv3d(16, 16, 3, 1, padding=1),
nn.PReLU(16),
)
self.encoder_stage2 = nn.Sequential(
nn.Conv3d(32, 32, 3, 1, padding=1),
nn.PReLU(32),
nn.Conv3d(32, 32, 3, 1, padding=1),
nn.PReLU(32),
nn.Conv3d(32, 32, 3, 1, padding=1),
nn.PReLU(32),
)
self.encoder_stage3 = nn.Sequential(
nn.Conv3d(64, 64, 3, 1, padding=1),
nn.PReLU(64),
nn.Conv3d(64, 64, 3, 1, padding=2, dilation=2),
nn.PReLU(64),
nn.Conv3d(64, 64, 3, 1, padding=4, dilation=4),
nn.PReLU(64),
)
self.encoder_stage4 = nn.Sequential(
nn.Conv3d(128, 128, 3, 1, padding=3, dilation=3),
nn.PReLU(128),
nn.Conv3d(128, 128, 3, 1, padding=4, dilation=4),
nn.PReLU(128),
nn.Conv3d(128, 128, 3, 1, padding=5, dilation=5),
nn.PReLU(128),
)
self.decoder_stage1 = nn.Sequential(
nn.Conv3d(128, 256, 3, 1, padding=1),
nn.PReLU(256),
nn.Conv3d(256, 256, 3, 1, padding=1),
nn.PReLU(256),
nn.Conv3d(256, 256, 3, 1, padding=1),
nn.PReLU(256),
)
self.decoder_stage2 = nn.Sequential(
nn.Conv3d(128 + 64, 128, 3, 1, padding=1),
nn.PReLU(128),
nn.Conv3d(128, 128, 3, 1, padding=1),
nn.PReLU(128),
nn.Conv3d(128, 128, 3, 1, padding=1),
nn.PReLU(128),
)
self.decoder_stage3 = nn.Sequential(
nn.Conv3d(64 + 32, 64, 3, 1, padding=1),
nn.PReLU(64),
nn.Conv3d(64, 64, 3, 1, padding=1),
nn.PReLU(64),
nn.Conv3d(64, 64, 3, 1, padding=1),
nn.PReLU(64),
)
self.decoder_stage4 = nn.Sequential(
nn.Conv3d(32 + 16, 32, 3, 1, padding=1),
nn.PReLU(32),
nn.Conv3d(32, 32, 3, 1, padding=1),
nn.PReLU(32),
)
self.down_conv1 = nn.Sequential(nn.Conv3d(16, 32, 2, 2), nn.PReLU(32))
self.down_conv2 = nn.Sequential(nn.Conv3d(32, 64, 2, 2), nn.PReLU(64))
self.down_conv3 = nn.Sequential(nn.Conv3d(64, 128, 2, 2),
nn.PReLU(128))
self.down_conv4 = nn.Sequential(nn.Conv3d(128, 256, 3, 1, padding=1),
nn.PReLU(256))
self.up_conv2 = nn.Sequential(nn.ConvTranspose3d(256, 128, 2, 2),
nn.PReLU(128))
self.up_conv3 = nn.Sequential(nn.ConvTranspose3d(128, 64, 2, 2),
nn.PReLU(64))
self.up_conv4 = nn.Sequential(nn.ConvTranspose3d(64, 32, 2, 2),
nn.PReLU(32))
# 最后大尺度下的映射（256*256），下面的尺度依次递减
self.map4 = nn.Sequential(nn.Conv3d(32, 1, 1, 1), nn.Sigmoid())
# 128*128 尺度下的映射
self.map3 = nn.Sequential(
nn.Conv3d(64, 1, 1, 1),
nn.Upsample(scale_factor=2, mode='trilinear'), nn.Sigmoid())
# 64*64 尺度下的映射
self.map2 = nn.Sequential(
nn.Conv3d(128, 1, 1, 1),
nn.Upsample(scale_factor=4, mode='trilinear'), nn.Sigmoid())
# 32*32 尺度下的映射
self.map1 = nn.Sequential(
nn.Conv3d(256, 1, 1, 1),
nn.Upsample(scale_factor=8, mode='trilinear'), nn.Sigmoid())
def __init__(self,
filters=[64, 128, 256, 512],
d_model=768,
input_shape=(512, 512, 512),
patch_size=(16, 16, 16),
skip_idx=[3, 6, 9, 12],
n_classes=2,
in_channels=1,
n_heads=8,
bn=True,
up_mode='deconv',
n_layers=12):
super(UNETR, self).__init__()
print("UNETR")
self.in_channels = in_channels
self.d_model = d_model
self.n_heads = n_heads
self.input_shape = input_shape
self.patch_size = patch_size
self.n_layers = n_layers
self.filters = filters
self.filters.reverse()
self.skip_idx = skip_idx
self.emb_size_reshape = [
int(i / j) for i, j in zip(self.input_shape, self.patch_size)
] + [np.prod(self.patch_size)]
self.emb_size_flat = [
np.prod(self.emb_size_reshape[:3]), self.emb_size_reshape[3]
]
print('self.emb_size_reshape', self.emb_size_reshape)
print('self.emb_size_flat', self.emb_size_flat)
# Encoders
self.lin = nn.Linear(self.emb_size_reshape[3], self.d_model)
self.ListTrans = []
for i in range(self.n_layers):
encoder_layer = nn.TransformerEncoderLayer(d_model=self.d_model,
nhead=self.n_heads)
self.ListTrans.append(nn.TransformerEncoder(encoder_layer, 1))
self.TransModuleList = nn.ModuleList(self.ListTrans)
# Skips
self.skip0 = UNetConv3D(self.in_channels, self.filters[3], bn=bn)
self.skip1 = UNETRSkip(self.d_model, self.filters[:3], bn=bn)
self.skip2 = UNETRSkip(self.d_model, self.filters[:2], bn=bn)
self.skip3 = UNETRSkip(self.d_model, self.filters[:1], bn=bn)
# Upsamplers
self.up_concat4 = nn.ConvTranspose3d(self.d_model,
self.filters[0],
2,
stride=2)
self.up_concat3 = nn.Sequential(*[
UNetConv3D(self.filters[0] * 2, self.filters[1], bn=bn),
nn.ConvTranspose3d(
self.filters[1], self.filters[1], (2, 2, 2), stride=(2, 2, 2))
])
self.up_concat2 = nn.Sequential(*[
UNetConv3D(self.filters[1] * 2, self.filters[2], bn=bn),
nn.ConvTranspose3d(
self.filters[2], self.filters[2], (2, 2, 2), stride=(2, 2, 2))
])
self.up_concat1 = nn.Sequential(*[
UNetConv3D(self.filters[2] * 2, self.filters[3], bn=bn),
nn.ConvTranspose3d(
self.filters[3], self.filters[3], (2, 2, 2), stride=(2, 2, 2))
])
# final conv (without any concat)
self.final = nn.Sequential(*[
UNetConv3D(self.filters[3] * 2, n_classes, bn=bn),
nn.Conv3d(n_classes, n_classes, kernel_size=1)
])
# initialise weights
for m in self.modules():
if isinstance(m, nn.ConvTranspose3d) or isinstance(
m, nn.Conv3d) or isinstance(m, nn.TransformerEncoderLayer):
init_weights(m, init_type='kaiming')
def __init__(self, n0, n1, n2, n3, p = 0.0, integrate = True):
super(Rec3, self).__init__()
self.block01 = nn.Sequential(
nn.Conv3d(n0, n1, kernel_size = 3, stride = 2, padding = 1),
nn.BatchNorm3d(n1),
nn.ReLU(inplace = True),
nn.Conv3d(n1, n1, kernel_size = 3, padding = 1),
nn.BatchNorm3d(n1))
self.block11 = nn.Sequential(
nn.Conv3d(n1, n1, kernel_size = 3, padding = 1),
nn.BatchNorm3d(n1),
nn.ReLU(inplace = True),
nn.Conv3d(n1, n1, kernel_size = 3, padding = 1),
nn.BatchNorm3d(n1))
self.block21 = nn.Sequential(
nn.ConvTranspose3d(n2, n1, kernel_size = 2, stride = 2),
nn.BatchNorm3d(n1),
nn.ReLU(inplace = True),
nn.Conv3d(n1, n1, kernel_size = 3, padding = 1),
nn.BatchNorm3d(n1))
self.block12 = nn.Sequential(
nn.Conv3d(n1, n2, kernel_size = 3, stride = 2, padding = 1),
nn.BatchNorm3d(n2),
nn.ReLU(inplace = True),
nn.Conv3d(n2, n2, kernel_size = 3, padding = 1),
nn.BatchNorm3d(n2))
self.block22 = nn.Sequential(
nn.Conv3d(n2, n2, kernel_size = 3, padding = 1),
nn.BatchNorm3d(n2),
nn.ReLU(inplace = True),
nn.Conv3d(n2, n2, kernel_size = 3, padding = 1),
nn.BatchNorm3d(n2))
self.block32 = nn.Sequential(
nn.ConvTranspose3d(n3, n2, kernel_size = 2, stride = 2),
nn.BatchNorm3d(n2),
nn.ReLU(inplace = True),
nn.Conv3d(n2, n2, kernel_size = 3, padding = 1),
nn.BatchNorm3d(n2))
self.block23 = nn.Sequential(
nn.Conv3d(n2, n3, kernel_size = 3, stride = 2, padding = 1),
nn.BatchNorm3d(n3),
nn.ReLU(inplace = True),
nn.Conv3d(n3, n3, kernel_size = 3, padding = 1),
nn.BatchNorm3d(n3))
self.block33 = nn.Sequential(
nn.Conv3d(n3, n3, kernel_size = 3, padding = 1),
nn.BatchNorm3d(n3),
nn.ReLU(inplace = True),
nn.Conv3d(n3, n3, kernel_size = 3, padding = 1),
nn.BatchNorm3d(n3))
self.relu = nn.ReLU(inplace = True)
self.p = p
self.integrate = integrate
def __init__(self,
in_channel,
channel,
n_res_block,
n_res_channel,
stride,
activation,
dense_layers_sizes,
is_bns,
is_dropouts,
final_activation=None,
drop_val=0.5,
is_bayesian=False,
random_node="output"):
super().__init__()
self.is_bayesian = is_bayesian
self.blocks1 = ResBlockDeconv3D(channel, n_res_channel)
self.blocks2 = ResBlockDeconv3D(channel, n_res_channel)
self.blocks3 = ResBlockDeconv3D(channel, n_res_channel)
self.blocks4 = ResBlockDeconv3D(channel, n_res_channel)
self.blocks5 = nn.ConvTranspose3d(channel, channel, 3, padding=1)
self.blocks6 = nn.ConvTranspose3d(channel,
channel,
4,
stride=2,
padding=1)
self.blocks7 = nn.ConvTranspose3d(channel,
channel,
4,
stride=2,
padding=1)
self.blocks8 = nn.ConvTranspose3d(channel,
channel,
4,
stride=2,
padding=1)
self.blocks9 = nn.ConvTranspose3d(channel,
channel,
4,
stride=2,
padding=1)
self.blocks10 = nn.ConvTranspose3d(channel,
channel,
4,
stride=2,
padding=1)
self.blocks11 = nn.ConvTranspose3d(channel,
channel,
4,
stride=2,
padding=1)
self.blocks12 = nn.ConvTranspose3d(channel,
channel,
4,
stride=2,
padding=1)
self.blocks13 = nn.ConvTranspose3d(channel,
channel,
4,
stride=2,
padding=1)
self.blocks14 = nn.ConvTranspose3d(channel,
channel,
4,
stride=2,
padding=1)
self.blocks15 = nn.ConvTranspose3d(channel,
channel,
4,
stride=2,
padding=1)
self.blocks16 = nn.ConvTranspose3d(channel,
channel,
4,
stride=2,
padding=1)
self.blocks17 = nn.ConvTranspose3d(channel,
channel,
4,
stride=2,
padding=1)
self.blocks18 = nn.ConvTranspose3d(channel,
channel,
4,
stride=2,
padding=1)
self.blocks19 = nn.ConvTranspose3d(channel,
channel,
4,
stride=2,
padding=1)
self.blocks20 = nn.ConvTranspose3d(channel,
channel,
4,
stride=2,
padding=1)
self.blocks21 = nn.ConvTranspose3d(channel,
channel,
4,
stride=2,
padding=1)
self.blocks22 = nn.ConvTranspose3d(channel,
channel // 2,
4,
stride=2,
padding=1)
self.blocks23 = nn.ConvTranspose3d(channel // 2,
in_channel,
4,
stride=2,
padding=1)
self.is_dropouts = is_dropouts
self.dropout = [[] for _ in dense_layers_sizes]
self.bns = [[] for _ in dense_layers_sizes]
self.linears = [[] for _ in dense_layers_sizes]
self.bn0 = torch.nn.BatchNorm3d(dense_layers_sizes[0])
self.is_bns = is_bns
for i in range(len(dense_layers_sizes) - 1):
self.linears[i] = torch.nn.Linear(
in_features=dense_layers_sizes[i],
out_features=dense_layers_sizes[i + 1]).to(device)
if self.is_bns[i] == 1:
self.bns[i] = torch.nn.BatchNorm3d(
dense_layers_sizes[i]).to(device)
else:
self.bns[i] = None
if self.is_dropouts[i] == 1:
self.dropout[i] = nn.Dropout(drop_val).to(device)
else:
self.dropout[i] = None
self.activation = activation
self.final_activation = final_activation
def __init__(self, opt):
super(PrimaryNet, self).__init__()
self._opt = opt
self.lossCriterion = nn.MSELoss()
repPad = nn.ReplicationPad3d((1, 1, 1, 1, 0, 0))
self.conv0 = nn.Sequential(
repPad, nn.Conv3d(1, 8, (1, 3, 3), stride=1, padding=0),
nn.InstanceNorm3d(8), nn.LeakyReLU(), repPad,
nn.Conv3d(8, 8, (1, 3, 3), stride=1, padding=0),
nn.InstanceNorm3d(8), nn.LeakyReLU())
self.convH0 = nn.Sequential(
repPad, nn.Conv3d(8, 8, (1, 3, 3), stride=1, padding=0),
nn.InstanceNorm3d(8), nn.LeakyReLU(), repPad,
nn.Conv3d(8, 8, (1, 3, 3), stride=1, padding=0),
nn.InstanceNorm3d(8), nn.LeakyReLU(), repPad,
nn.Conv3d(8, 8, (1, 3, 3), stride=1, padding=0),
nn.InstanceNorm3d(8), nn.LeakyReLU(), repPad,
nn.Conv3d(8, 8, (1, 3, 3), stride=1, padding=0),
nn.InstanceNorm3d(8), nn.LeakyReLU())
self.convL0 = nn.Sequential(
repPad,
nn.Conv3d(8, 16, (1, 3, 3), stride=(1, 2, 2), padding=0),
nn.InstanceNorm3d(16),
nn.LeakyReLU(), #192,48
repPad,
nn.Conv3d(16, 16, (1, 3, 3), stride=(1, 1, 1), padding=0),
nn.InstanceNorm3d(16),
nn.LeakyReLU(),
repPad,
nn.Conv3d(16, 32, (1, 3, 3), stride=(1, 2, 2), padding=0),
nn.InstanceNorm3d(32),
nn.LeakyReLU(), #96,24
repPad,
nn.Conv3d(32, 32, (1, 3, 3), stride=(1, 1, 1), padding=0),
nn.InstanceNorm3d(32),
nn.LeakyReLU(),
repPad,
nn.Conv3d(32, 48, (1, 3, 3), stride=(1, 2, 2), padding=0),
nn.InstanceNorm3d(48),
nn.LeakyReLU()) #48,12
self.convL0Up = nn.Sequential(
nn.ConvTranspose3d(48,
32, (1, 7, 7),
stride=(1, 4, 4),
padding=(0, 3, 3),
output_padding=(0, 3, 3)),
nn.ConvTranspose3d(32,
16, (1, 5, 5),
stride=(1, 2, 2),
padding=(0, 2, 2),
output_padding=(0, 1, 1)))
self.convH1 = nn.Sequential(
repPad, nn.Conv3d(24, 12, (1, 3, 3), stride=1, padding=0),
nn.InstanceNorm3d(12), nn.LeakyReLU(), repPad,
nn.Conv3d(12, 12, (1, 3, 3), stride=1, padding=0),
nn.InstanceNorm3d(12), nn.LeakyReLU(), repPad,
nn.Conv3d(12, 12, (1, 3, 3), stride=1, padding=0),
nn.InstanceNorm3d(12), nn.LeakyReLU(), repPad,
nn.Conv3d(12, 12, (1, 3, 3), stride=1, padding=0),
nn.InstanceNorm3d(12), nn.LeakyReLU())
self.convL1 = nn.Sequential(
repPad, nn.Conv3d(48, 48, (1, 3, 3), stride=(1, 1, 1), padding=0),
nn.InstanceNorm3d(48), nn.LeakyReLU(), repPad,
nn.Conv3d(48, 48, (1, 3, 3), stride=(1, 1, 1), padding=0),
nn.InstanceNorm3d(48), nn.LeakyReLU(), repPad,
nn.Conv3d(48, 48, (1, 3, 3), stride=(1, 1, 1), padding=0),
nn.InstanceNorm3d(48), nn.LeakyReLU(), repPad,
nn.Conv3d(48, 48, (1, 3, 3), stride=(1, 1, 1), padding=0),
nn.InstanceNorm3d(48), nn.LeakyReLU())
self.convL1Up = nn.Sequential(
nn.ConvTranspose3d(48,
32, (1, 7, 7),
stride=(1, 4, 4),
padding=(0, 3, 3),
output_padding=(0, 3, 3)),
nn.ConvTranspose3d(32,
16, (1, 5, 5),
stride=(1, 2, 2),
padding=(0, 2, 2),
output_padding=(0, 1, 1)))
self.convH2 = nn.Sequential(
repPad, nn.Conv3d(28, 12, (1, 3, 3), stride=1, padding=0),
nn.InstanceNorm3d(12), nn.LeakyReLU(), repPad,
nn.Conv3d(12, 12, (1, 3, 3), stride=1, padding=0),
nn.InstanceNorm3d(12), nn.LeakyReLU(), repPad,
nn.Conv3d(12, 12, (1, 3, 3), stride=1, padding=0),
nn.InstanceNorm3d(12), nn.LeakyReLU(), repPad,
nn.Conv3d(12, 12, (1, 3, 3), stride=1, padding=0),
nn.InstanceNorm3d(12), nn.LeakyReLU())
self.convL2 = nn.Sequential(
repPad, nn.Conv3d(48, 48, (1, 3, 3), stride=(1, 1, 1), padding=0),
nn.InstanceNorm3d(48), nn.LeakyReLU(), repPad,
nn.Conv3d(48, 48, (1, 3, 3), stride=(1, 1, 1), padding=0),
nn.InstanceNorm3d(48), nn.LeakyReLU(), repPad,
nn.Conv3d(48, 48, (1, 3, 3), stride=(1, 1, 1), padding=0),
nn.InstanceNorm3d(48), nn.LeakyReLU(), repPad,
nn.Conv3d(48, 48, (1, 3, 3), stride=(1, 1, 1), padding=0),
nn.InstanceNorm3d(48), nn.LeakyReLU())
self.convL2Up = nn.Sequential(
nn.ConvTranspose3d(48,
32, (1, 7, 7),
stride=(1, 4, 4),
padding=(0, 3, 3),
output_padding=(0, 3, 3)),
nn.ConvTranspose3d(32,
16, (1, 5, 5),
stride=(1, 2, 2),
padding=(0, 2, 2),
output_padding=(0, 1, 1)))
self.convProj = nn.Sequential(
repPad, nn.Conv3d(28, 1, (1, 3, 3), stride=1, padding=0))
def __init__(self):
super(EDCNN, self).__init__()
self.feature_maps = 32
self.kernel_size = (3, 5, 5)
self.stride = 1
#Contracting path:
self.enc_1 = nn.Sequential(
nn.Conv3d(1,
self.feature_maps,
kernel_size=self.kernel_size,
stride=self.stride), nn.ReLU())
self.enc_2 = nn.Sequential(
nn.Conv3d(self.feature_maps,
self.feature_maps,
kernel_size=self.kernel_size,
stride=self.stride), nn.ReLU())
self.enc_3 = nn.Sequential(
nn.Conv3d(self.feature_maps,
self.feature_maps,
kernel_size=self.kernel_size,
stride=self.stride), nn.ReLU())
self.enc_4 = nn.Sequential(
nn.Conv3d(self.feature_maps,
self.feature_maps,
kernel_size=self.kernel_size,
stride=self.stride), nn.ReLU())
#Expansive path
self.dec_1 = nn.Sequential(
nn.ConvTranspose3d(self.feature_maps,
self.feature_maps,
kernel_size=self.kernel_size,
stride=self.stride), nn.ReLU())
self.dec_2 = nn.Sequential(
nn.ConvTranspose3d(self.feature_maps,
self.feature_maps,
kernel_size=self.kernel_size,
stride=self.stride), nn.ReLU())
self.dec_3 = nn.Sequential(
nn.ConvTranspose3d(self.feature_maps,
self.feature_maps,
kernel_size=self.kernel_size,
stride=self.stride), nn.ReLU())
self.dec_1 = nn.Sequential(
nn.ConvTranspose3d(self.feature_maps,
self.feature_maps,
kernel_size=self.kernel_size,
stride=self.stride), nn.ReLU())
self.dec_4 = nn.ConvTranspose3d(self.feature_maps,
1,
kernel_size=self.kernel_size,
stride=self.stride)
def __init__(self):
super(Net, self).__init__()
self.preBlock = nn.Sequential(
nn.Conv3d(1, 24, kernel_size=3, padding=1), nn.BatchNorm3d(24),
nn.ReLU(), nn.Conv3d(24, 24, kernel_size=3, padding=1),
nn.BatchNorm3d(24), nn.ReLU())
# 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]
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))
self.nodule_output = nn.Sequential(
nn.Conv3d(self.featureNum_back[0], 64, kernel_size=1),
nn.ReLU(),
#nn.Dropout3d(p = 0.3),
nn.Conv3d(64, len(config['anchors']), kernel_size=1))
self.regress_output = nn.Sequential(
nn.Conv3d(self.featureNum_back[0], 64, kernel_size=1),
nn.ReLU(),
#nn.Dropout3d(p = 0.3),
nn.Conv3d(64, 4 * len(config['anchors']), kernel_size=1))
focal_bias = -math.log((1.0 - 0.01) / 0.01)
self._modules['nodule_output'][2].bias.data.fill_(focal_bias)
def __init__(self,
input_shape,
in_channels,
kernel_size=3,
padding=1,
flow_multiplier=1.,
nb_channels=16):
""" Init class.
Parameters
----------
input_shape: uplet
the tensor data shape (X, Y, Z).
in_channels: int
number of channels in the input tensor.
kernel_size: int, default 3
the convolution kernels size (odd number).
padding: int, default 1
the padding size, recommended (kernel_size - 1) / 2
flow_multiplier: foat, default 1
weight the flow field by this factor.
nb_channels: int, default 16
the number of channels after the first convolution.
"""
# Inheritance
nn.Module.__init__(self)
# Class parameters
self.input_shape = input_shape
self.in_channels = in_channels
self.kernel_size = kernel_size
self.padding = padding
self.flow_multiplier = flow_multiplier
self.shapes = self._downsample_shape(input_shape,
nb_iterations=6,
scale_factor=2)
self.nb_channels = nb_channels
# Use strided 3D convolution to progressively downsample the image,
# and then use deconvolution (transposed convolution) to recover
# spatial resolution. As suggested in U-Net, skip connections between
# the convolutional layers and the deconvolutional layers are added
# to help refining dense prediction. The network will output the dense
# flow field, a volume feature map with 3 channels (X, Y, Z
# displacements) of the same size as the input.
out_channels = nb_channels
for idx in range(1, 3):
ops = self._conv(in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
stride=2,
padding=1,
bias=True,
negative_slope=0.1)
setattr(self, "down{0}".format(idx), ops)
in_channels = out_channels
out_channels *= 2
for idx in range(3, 7):
ops = self._double_conv(in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
stride=2,
padding=1,
bias=True,
negative_slope=0.1)
setattr(self, "down{0}".format(idx), ops)
in_channels = out_channels
out_channels *= 2
out_channels = in_channels // 2
for idx in range(5, 0, -1):
if idx < 5:
in_channels = in_channels * 2 + 3
pred_ops = self._prediction(in_channels=in_channels)
ops = self._upconv(in_channels=in_channels,
out_channels=out_channels,
kernel_size=4,
stride=2,
padding=2,
groups=1,
negative_slope=0.1)
setattr(self, "pred{0}".format(idx + 1), pred_ops)
setattr(self, "up{0}".format(idx), ops)
in_channels = out_channels
out_channels = out_channels // 2
in_channels = in_channels * 2 + 3
self.pred1 = nn.ConvTranspose3d(in_channels=in_channels,
out_channels=3,
kernel_size=4,
stride=2,
padding=1,
groups=1)
# Finally warp the moving image.
self.spatial_transform = SpatialTransformer(input_shape)
# Init weights
@torch.no_grad()
def weights_init(module):
if isinstance(module, nn.Conv3d):
logger.debug("Init weights of {0}...".format(module))
torch.nn.init.xavier_uniform_(module.weight)
torch.nn.init.constant_(module.bias, 0)
self.apply(weights_init)
def __init__(self, in_channels, n_classes):
super(VoxResNet_AG, self).__init__()
self.conv1 = nn.Conv3d(in_channels=in_channels,
out_channels=32,
kernel_size=3,
padding=1)
self.bn1 = nn.BatchNorm3d(32)
self.act1 = ActFunc('ReLU')
self.conv2 = nn.Conv3d(in_channels=32,
out_channels=32,
kernel_size=3,
padding=1)
self.bn2 = nn.BatchNorm3d(32)
self.act2 = ActFunc('ReLU')
self.conv3 = nn.Conv3d(in_channels=32,
out_channels=64,
kernel_size=3,
padding=1,
stride=2)
self.mod1 = VoxResModule(in_channels=64)
self.mod2 = VoxResModule(in_channels=64)
self.bn3 = nn.BatchNorm3d(64)
self.act3 = ActFunc('ReLU')
self.conv4 = nn.Conv3d(in_channels=64,
out_channels=64,
kernel_size=3,
padding=1,
stride=2)
self.mod3 = VoxResModule(in_channels=64)
self.mod4 = VoxResModule(in_channels=64)
self.bn4 = nn.BatchNorm3d(64)
self.act4 = ActFunc('ReLU')
self.conv5 = nn.Conv3d(in_channels=64,
out_channels=64,
kernel_size=3,
padding=1,
stride=2)
self.mod5 = VoxResModule(in_channels=64)
self.mod6 = VoxResModule(in_channels=64)
self.gs = GatingSignal(in_channels=64, out_channels=64)
self.attention1 = AttentionGate(in_channels=32,
gating_channels=64,
inter_channels=16)
self.attention2 = AttentionGate(in_channels=64,
gating_channels=64,
inter_channels=32)
self.attention3 = AttentionGate(in_channels=64,
gating_channels=64,
inter_channels=32)
self.attention4 = AttentionGate(in_channels=64,
gating_channels=64,
inter_channels=32)
# Deconvolution layers
self.deconv1 = nn.ConvTranspose3d(in_channels=32,
out_channels=n_classes,
kernel_size=3,
stride=1,
padding=1)
self.deconv2 = nn.ConvTranspose3d(in_channels=64,
out_channels=n_classes,
kernel_size=2,
stride=2)
self.deconv3 = nn.ConvTranspose3d(in_channels=64,
out_channels=n_classes,
kernel_size=4,
stride=4)
self.deconv4 = nn.ConvTranspose3d(in_channels=64,
out_channels=n_classes,
kernel_size=8,
stride=8)
self.softmax = F.softmax
def __init__(self,
in_channels=16,
input_data_channels=1,
output_label_channels=1,
dropout_rate=0):
super().__init__()
self.actual_epoch = 0
self.actual_step = 0
out_channels = in_channels
# downconv1 y=>x, x=>2x
double_out_channels = out_channels * 2
self.dconv_down1 = nn.Sequential(
double_conv(input_data_channels, out_channels,
double_out_channels), nn.Dropout(dropout_rate))
# sSE
self.dconv_atten1 = attention_block(in_channels=double_out_channels)
out_channels *= 2
self.pool1 = nn.MaxPool3d(2, stride=2, ceil_mode=True) # 1/2
self.atten_pool1 = nn.MaxPool3d(2, stride=2, ceil_mode=True) # 1/2
# downconv2 2x=>2x, 2x=>4x
double_out_channels = out_channels * 2
self.dconv_down2 = nn.Sequential(
down_double_conv(out_channels, double_out_channels),
nn.Dropout(dropout_rate))
# sSE
self.dconv_atten2 = attention_block(in_channels=double_out_channels)
self.dconv_merge_atten2 = attention_block(in_channels=2)
out_channels *= 2
self.pool2 = nn.MaxPool3d(2, stride=2, ceil_mode=True) # 1/4
self.atten_pool2 = nn.MaxPool3d(2, stride=2, ceil_mode=True) # 1/4
# downconv3 4x=>4x, 4x=>8x
double_out_channels = out_channels * 2
self.dconv_down3 = nn.Sequential(
down_double_conv(out_channels, double_out_channels),
nn.Dropout(dropout_rate))
# sSE
self.dconv_atten3 = attention_block(in_channels=double_out_channels)
self.dconv_merge_atten3 = attention_block(in_channels=2)
out_channels = double_out_channels
self.pool3 = nn.MaxPool3d(2, stride=2, ceil_mode=True) # 1/8
self.atten_pool3 = nn.MaxPool3d(2, stride=2, ceil_mode=True) # 1/4
# middle 8x=>8x, 8x=>16x
double_out_channels = out_channels * 2
self.dconv_middle = nn.Sequential(
down_double_conv(out_channels, double_out_channels),
nn.Dropout(dropout_rate))
# sSE
self.middle_atten = attention_block(in_channels=double_out_channels)
self.middle_merge_atten = attention_block(in_channels=2)
out_channels = double_out_channels
logging.debug(
f'UNet Architecture v2: max output channels {out_channels}')
# upconv1 16x+8x=>8x, 8x=>8x
half_out_channels = out_channels // 2
self.up1 = nn.ConvTranspose3d(out_channels,
out_channels,
2,
stride=2,
bias=False)
self.dconv_up1 = up_double_conv(out_channels + half_out_channels,
half_out_channels)
# sSE
self.atten_up1 = nn.Upsample(scale_factor=2)
self.atten_dconv_up1 = attention_block(in_channels=half_out_channels)
self.atten_dconv_merge_up1 = attention_block(in_channels=2)
out_channels = half_out_channels
# upconv2 8x+4x=>4x, 4x=>4x
half_out_channels = out_channels // 2
self.up2 = nn.ConvTranspose3d(out_channels,
out_channels,
2,
stride=2,
bias=False)
self.dconv_up2 = up_double_conv(out_channels + half_out_channels,
half_out_channels)
# sSE
self.atten_up2 = nn.Upsample(scale_factor=2)
self.atten_dconv_up2 = attention_block(in_channels=half_out_channels)
self.atten_dconv_merge_up2 = attention_block(in_channels=2)
out_channels = half_out_channels
# upconv3 4x+2x=>2x, 2x=>2x
half_out_channels = out_channels // 2
self.up3 = nn.ConvTranspose3d(out_channels,
out_channels,
2,
stride=2,
bias=False)
self.dconv_up3 = up_double_conv(out_channels + half_out_channels,
half_out_channels)
# sSE
self.atten_up3 = nn.Upsample(scale_factor=2)
self.atten_dconv_up3 = attention_block(in_channels=half_out_channels)
self.atten_dconv_merge_up3 = attention_block(in_channels=2)
out_channels = half_out_channels
# final conv 2x=>z
self.final = nn.Conv3d(out_channels,
output_label_channels,
kernel_size=1)
self.sigmoid = nn.Sigmoid()
def __init__(self, nf_in, nf_per_level, nf_out, use_skip_sparse,
use_skip_dense, input_volume_size):
nn.Module.__init__(self)
assert (type(nf_per_level) is list)
data_dim = 3
self.use_skip_sparse = use_skip_sparse
self.use_skip_dense = use_skip_dense
#self.use_bias = True
self.use_bias = False
modules = []
volume_sizes = [(np.array(input_volume_size) // (k + 1)).tolist()
for k in range(len(nf_per_level))]
for level in range(len(nf_per_level)):
nf_in = nf_in if level == 0 else nf_per_level[level - 1]
input_sparsetensor = level > 0
return_sparsetensor = (level < len(nf_per_level) - 1)
modules.append(
SparseEncoderLayer(nf_in, nf_per_level[level],
input_sparsetensor, return_sparsetensor,
volume_sizes[level]))
self.process_sparse = nn.Sequential(*modules)
nf = nf_per_level[-1]
# 16 -> 8
nf0 = nf * 3 // 2
self.encode_dense0 = nn.Sequential(
nn.Conv3d(nf,
nf0,
kernel_size=4,
stride=2,
padding=1,
bias=self.use_bias), nn.BatchNorm3d(nf0), nn.ReLU(True))
# 8 -> 4
nf1 = nf * 2
self.encode_dense1 = nn.Sequential(
nn.Conv3d(nf0,
nf1,
kernel_size=4,
stride=2,
padding=1,
bias=self.use_bias), nn.BatchNorm3d(nf1), nn.ReLU(True))
# 4 -> 4
nf2 = nf1
self.bottleneck_dense2 = nn.Sequential(
nn.Conv3d(nf1, nf2, kernel_size=1, bias=self.use_bias),
nn.BatchNorm3d(nf2), nn.ReLU(True))
# 4 -> 8
nf3 = nf2 if not self.use_skip_dense else nf1 + nf2
nf4 = nf3 // 2
self.decode_dense3 = nn.Sequential(
nn.ConvTranspose3d(nf3,
nf4,
kernel_size=4,
stride=2,
padding=1,
bias=self.use_bias), nn.BatchNorm3d(nf4),
nn.ReLU(True))
# 8 -> 16
if self.use_skip_dense:
nf4 += nf0
nf5 = nf4 // 2
self.decode_dense4 = nn.Sequential(
nn.ConvTranspose3d(nf4,
nf5,
kernel_size=4,
stride=2,
padding=1,
bias=self.use_bias), nn.BatchNorm3d(nf5),
nn.ReLU(True))
self.final = nn.Sequential(
nn.Conv3d(nf5, nf_out, kernel_size=1, bias=self.use_bias),
nn.BatchNorm3d(nf_out), nn.ReLU(True))
# occ prediction
self.occpred = nn.Sequential(
nn.Conv3d(nf_out, 1, kernel_size=1, bias=self.use_bias))
self.sdfpred = nn.Sequential(
nn.Conv3d(nf_out, 1, kernel_size=1, bias=self.use_bias))
# debug stats
params_encodesparse = count_num_model_params(self.process_sparse)
params_encodedense = count_num_model_params(
self.encode_dense0) + count_num_model_params(
self.encode_dense0) + count_num_model_params(
self.encode_dense1) + count_num_model_params(
self.bottleneck_dense2)
params_decodedense = count_num_model_params(
self.decode_dense3) + count_num_model_params(
self.decode_dense4) + count_num_model_params(
self.final) + count_num_model_params(self.occpred)
print('[TSDFEncoder] params encode sparse', params_encodesparse)
print('[TSDFEncoder] params encode dense', params_encodedense)
print('[TSDFEncoder] params decode dense', params_decodedense)
def __init__(self):
super(HDR, self).__init__()
self.conv1 = nn.Conv3d(3, 16, 3,
padding=(1, 1, 1)) #[b, 16, 3, 180, 320]
self.bn1 = nn.BatchNorm3d(16)
self.conv2 = nn.Conv3d(16, 16, 3,
padding=(1, 1, 1)) #[b, 16, 3, 180, 320]
self.bn2 = nn.BatchNorm3d(16)
self.mpool1 = nn.MaxPool3d((1, 2, 2),
return_indices=True) #[b, 16, 3, 90, 160]
self.conv3 = nn.Conv3d(16, 32, 3,
padding=(1, 1, 1)) #[b, 32, 3, 90, 160]
self.bn3 = nn.BatchNorm3d(32)
self.conv4 = nn.Conv3d(32, 32, 3,
padding=(1, 1, 1)) #[b, 32, 3, 90, 160]
self.bn4 = nn.BatchNorm3d(32)
self.mpool2 = nn.MaxPool3d((1, 2, 2),
return_indices=True) #[b, 32, 3, 45, 80]
self.conv5 = nn.Conv3d(32, 64, 3,
padding=(0, 1, 1)) #[b, 64, 1, 45, 80]
self.bn5 = nn.BatchNorm3d(64)
self.conv6 = nn.Conv3d(64, 64, 3,
padding=(1, 1, 1)) #[b, 64, 1, 45, 80]
self.bn6 = nn.BatchNorm3d(64)
self.conv7 = nn.Conv3d(64, 64, 3,
padding=(1, 1, 1)) #[b, 64, 1, 45, 80]
self.bn7 = nn.BatchNorm3d(64)
self.mpool3 = nn.MaxPool3d((1, 2, 2),
return_indices=True) #[b, 64, 1, 22, 40]
self.conv8 = nn.Conv3d(64, 128, 3,
padding=(1, 1, 1)) #[b, 128, 1, 22, 40]
self.bn8 = nn.BatchNorm3d(128)
self.conv9 = nn.Conv3d(128, 128, (3, 3, 3),
padding=(1, 1, 1)) #[b, 128, 1, 22, 40]
self.bn9 = nn.BatchNorm3d(128)
self.conv10 = nn.Conv3d(128, 128, 3,
padding=(1, 1, 1)) #[b, 128, 1, 22, 40]
self.bn10 = nn.BatchNorm3d(128)
self.mpool4 = nn.MaxPool3d((1, 2, 2),
return_indices=True) #[b, 128, 1, 11, 20]
self.conv11 = nn.Conv3d(128, 256, 3,
padding=(1, 1, 1)) #[b, 256, 1, 11, 20]
self.bn11 = nn.BatchNorm3d(256)
self.conv12 = nn.Conv3d(256, 256, 3,
padding=(1, 1, 1)) #[b, 256, 1, 11, 20]
self.bn12 = nn.BatchNorm3d(256)
self.conv13 = nn.Conv3d(256, 256, 3,
padding=(1, 1, 1)) #[b, 256, 1, 11, 20]
self.bn13 = nn.BatchNorm3d(256)
self.mpool5 = nn.MaxPool3d((1, 2, 2),
return_indices=True) #[b, 256, 1, 5, 10]
#latent-transform
self.conv14 = nn.Conv3d(256, 256, 3,
padding=(1, 1, 1)) #[b, 256, 1, 5, 10]
self.bn14 = nn.BatchNorm3d(256)
#decoder
self.unpool5 = nn.MaxUnpool3d((1, 2, 2)) #[b, 256, 1, 10, 20]
self.deconv13 = nn.ConvTranspose3d(256, 256, 3,
padding=(1, 1,
1)) #[b, 256, 1, 11, 20]
self.bn15 = nn.BatchNorm3d(256)
# skip cat
self.one_conv13 = nn.Conv3d(512, 256, 1) #[b, 256, 1, 11, 20]
self.deconv12 = nn.ConvTranspose3d(256, 256, 3,
padding=(1, 1,
1)) #[b, 256, 1, 11, 20]
self.bn16 = nn.BatchNorm3d(256)
self.deconv11 = nn.ConvTranspose3d(256, 128, 3,
padding=(1, 1,
1)) #[b, 128, 1, 11, 20]
self.bn17 = nn.BatchNorm3d(128)
self.unpool4 = nn.MaxUnpool3d((1, 2, 2)) #[b, 128, 1, 22, 40]
self.deconv10 = nn.ConvTranspose3d(128, 128, 3,
padding=(1, 1,
1)) #[b, 128, 1, 22, 40]
self.bn18 = nn.BatchNorm3d(128)
# skip cat
self.one_conv10 = nn.Conv3d(256, 128, 1) #[b, 128, 1, 22, 40]
self.deconv9 = nn.ConvTranspose3d(128,
128, (3, 3, 3),
padding=(1, 1, 1))
self.bn19 = nn.BatchNorm3d(128)
self.deconv8 = nn.ConvTranspose3d(128, 64, 3,
padding=(1, 1,
1)) #[b, 64, 1, 22, 40]
self.bn20 = nn.BatchNorm3d(64)
self.unpool3 = nn.MaxUnpool3d((1, 2, 2)) #[b, 64, 3, 45, 80]
self.deconv7 = nn.ConvTranspose3d(64, 64, 3,
padding=(1, 1,
1)) #[b, 64, 1, 45, 80]
self.bn21 = nn.BatchNorm3d(64)
# skip cat
self.one_conv7 = nn.Conv3d(128, 64, 1)
self.deconv6 = nn.ConvTranspose3d(64, 64, 3,
padding=(1, 1,
1)) #[b, 64, 1, 45, 80]
self.bn22 = nn.BatchNorm3d(64)
self.deconv5 = nn.ConvTranspose3d(64, 32, 3,
padding=(0, 1,
1)) #[b, 32, 3, 45, 80]
self.bn23 = nn.BatchNorm3d(32)
self.unpool2 = nn.MaxUnpool3d((1, 2, 2)) #[b, 32, 1, 90, 160]
self.deconv4 = nn.ConvTranspose3d(32, 32, 3,
padding=(1, 1,
1)) #[b, 32, 3, 90, 160]
self.bn24 = nn.BatchNorm3d(32)
self.one_conv4 = nn.Conv3d(64, 32, 1)
self.deconv3 = nn.ConvTranspose3d(32, 16, 3,
padding=(1, 1,
1)) #[b, 16, 3, 90, 160]
self.bn25 = nn.BatchNorm3d(16)
self.unpool1 = nn.MaxUnpool3d((1, 2, 2)) #[b, 16, 3, 180, 320]
self.deconv2 = nn.ConvTranspose3d(16, 16, 3,
padding=(1, 1,
1)) #[b, 16, 1, 180, 320]
self.bn26 = nn.BatchNorm3d(16)
self.one_conv2 = nn.Conv3d(32, 16, 1) #[b, 16, 3, 180, 320]
self.deconv1 = nn.ConvTranspose3d(16, 3, 3,
padding=(1, 1,
1)) #[b, 3, 3, 180, 320]
self.bn27 = nn.BatchNorm3d(3)
self.one_convop = nn.Conv3d(6, 3, (3, 1, 1))
# self.deconv0 = nn.ConvTranspose3d(3, 3, 3, padding=(2,1,1)) #[b, 3, 1, 180, 320]
# self.one_convop2 = nn.Conv3d(3, 3, (3,1,1), padding=(1,0,0))
self._initialize_weights()
def __init__(self, in_channels, out_channels): # in = out 5 4 4 = 128 96 96
super(Up, self).__init__()
# self.up = nn.Upsample(scale_factor=2, mode='trilinear', align_corners=True)
self.up = nn.ConvTranspose3d(in_channels, out_channels, 2, 2)
self.Att = Attention_block(F_g=out_channels, F_x=out_channels, F_int=in_channels) # 自己添加 ???
self.conv_block = nn.Conv3d(out_channels + out_channels, out_channels, 2, 2)
save_data_and_model("pool_conv_3d", input, model)
class Clip(nn.Module):
def __init__(self):
super(Clip, self).__init__()
def forward(self, x):
return torch.clamp(x, -0.1, 0.2)
model = Clip()
input = Variable(torch.rand(1, 10, 2, 2))
save_data_and_model('clip', input, model)
input = Variable(torch.randn(1, 3, 6, 6, 6))
deconv = nn.ConvTranspose3d(3, 3, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(0, 0, 0), bias=False)
save_data_and_model("deconv3d", input, deconv)
input = Variable(torch.randn(1, 3, 5, 4, 4))
deconv = nn.ConvTranspose3d(3, 5, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(0, 0, 0), bias=True)
save_data_and_model("deconv3d_bias", input, deconv)
input = Variable(torch.randn(1, 3, 5, 5, 5))
deconv = nn.ConvTranspose3d(3, 2, kernel_size=(4, 3, 3), stride=(1, 1, 1), padding=(1, 0, 1), bias=True)
save_data_and_model("deconv3d_pad", input, deconv)
input = Variable(torch.randn(1, 3, 4, 5, 3))
deconv = nn.ConvTranspose3d(3, 5, kernel_size=(2, 3, 1), stride=(2, 2, 2), padding=(1, 2, 1), output_padding=1, bias=True)
save_data_and_model("deconv3d_adjpad", input, deconv)
input = np.random.rand(1, 3, 4, 2)
def __init__(self, args):
super(PSMNet, self).__init__()
self.maxdisp = args.maxdisp
self.planes = args.planes
inplanes = self.planes * 2
self.feature_extraction = feature_extraction(args)
if args.shuffle:
block3d = ResBlock3DShuffle
else:
block3d = ResBlock3D
self.fuse = nn.Sequential(
nn.Conv3d(inplanes,
self.planes,
kernel_size=3,
stride=1,
padding=1,
dilation=1,
bias=False),
nn.BatchNorm3d(self.planes),
nn.ReLU(inplace=True),
)
self.unet_conv = nn.ModuleList()
inplanes = self.planes
for i in range(3):
outplanes = inplanes * 2
self.unet_conv.append(
nn.Sequential(
nn.Conv3d(inplanes,
outplanes,
kernel_size=3,
stride=2,
padding=3,
dilation=3,
bias=False),
nn.BatchNorm3d(outplanes),
nn.ReLU(inplace=True),
))
inplanes = outplanes
self.unet_dconv = nn.ModuleList()
self.classifiers = nn.ModuleList()
for i in range(3):
outplanes = inplanes // 2
self.unet_dconv.append(
nn.Sequential(
nn.ConvTranspose3d(inplanes,
outplanes,
kernel_size=3,
stride=2,
padding=3,
output_padding=1,
dilation=3,
bias=False),
nn.BatchNorm3d(outplanes),
nn.ReLU(inplace=True),
))
self.classifiers.append(
nn.Sequential(
nn.ConvTranspose3d(outplanes,
1,
kernel_size=7,
stride=4,
padding=3,
output_padding=3,
dilation=1,
bias=False)))
inplanes = outplanes
self.disparityregression = disparityregression(self.maxdisp)
for m in self.modules():
if isinstance(m, nn.Conv2d):
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
m.weight.data.normal_(0, math.sqrt(2. / n))
elif isinstance(m, nn.Conv3d):
n = m.kernel_size[0] * m.kernel_size[1] * m.kernel_size[
2] * m.out_channels
m.weight.data.normal_(0, math.sqrt(2. / n))
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
elif isinstance(m, nn.BatchNorm3d):
m.weight.data.fill_(1)
m.bias.data.zero_()
elif isinstance(m, nn.Linear):
m.bias.data.zero_()
def __init__(self,
block,
layers,
sample_input_D,
sample_input_H,
sample_input_W,
num_seg_classes,
shortcut_type='B',
no_cuda=False):
self.inplanes = 64
self.no_cuda = no_cuda
super(ResNet, self).__init__()
self.conv1 = nn.Conv3d(1,
64,
kernel_size=7,
stride=(2, 2, 2),
padding=(3, 3, 3),
bias=False)
self.bn1 = nn.BatchNorm3d(64)
self.relu = nn.ReLU(inplace=True)
self.maxpool = nn.MaxPool3d(kernel_size=(3, 3, 3), stride=2, padding=1)
self.layer1 = self._make_layer(block, 64, layers[0], shortcut_type)
self.layer2 = self._make_layer(block,
128,
layers[1],
shortcut_type,
stride=2)
self.layer3 = self._make_layer(block,
256,
layers[2],
shortcut_type,
stride=1,
dilation=2)
self.layer4 = self._make_layer(block,
512,
layers[3],
shortcut_type,
stride=1,
dilation=4)
self.conv_seg = nn.Sequential(
nn.ConvTranspose3d(512 * block.expansion, 512, 2, stride=2),
nn.BatchNorm3d(512), nn.ReLU(inplace=True),
nn.ConvTranspose3d(512, 256, 2, stride=2), nn.BatchNorm3d(256),
nn.ReLU(inplace=True), nn.ConvTranspose3d(256, 128, 2, stride=2),
nn.BatchNorm3d(128), nn.ReLU(inplace=True),
nn.Conv3d(128,
32,
kernel_size=3,
stride=(1, 1, 1),
padding=(1, 1, 1),
bias=False), nn.BatchNorm3d(32), nn.ReLU(inplace=True),
nn.Conv3d(32,
num_seg_classes,
kernel_size=1,
stride=(1, 1, 1),
bias=False))
for m in self.modules():
if isinstance(m, nn.Conv3d):
m.weight = nn.init.kaiming_normal_(m.weight, mode='fan_out')
elif isinstance(m, nn.BatchNorm3d):
m.weight.data.fill_(1)
m.bias.data.zero_()
def __init__(self,
nf_in,
nf_per_level,
nf_out,
use_skip_sparse=True,
use_skip_dense=True):
nn.Module.__init__(self)
assert (type(nf_per_level) is list)
data_dim = 3
self.use_skip_sparse = use_skip_sparse
self.use_skip_dense = use_skip_dense
self.use_aspp = False
#self.use_bias = True
self.use_bias = False
modules = []
for level in range(len(nf_per_level)):
nf_in = nf_in if level == 0 else nf_per_level[level - 1]
modules.append(PreEncoderLayer(nf_in, nf_per_level[level]))
self.process_sparse = nn.Sequential(*modules)
# print("TSDF SPARSE: ", self.process_sparse)
nf = nf_per_level[-1]
# 16 -> 8
nf0 = nf * 3 // 2
self.encode_dense0 = nn.Sequential(
nn.Conv3d(nf,
nf0,
kernel_size=4,
stride=2,
padding=1,
bias=self.use_bias), nn.BatchNorm3d(nf0), nn.ReLU(True))
nf1 = nf * 2
# 8 -> 4
self.encode_dense1 = nn.Sequential(
nn.Conv3d(nf0,
nf1,
kernel_size=4,
stride=2,
padding=1,
bias=self.use_bias), nn.BatchNorm3d(nf1), nn.ReLU(True))
# 4 -> 4
nf2 = nf1
self.bottleneck_dense2 = nn.Sequential(
nn.Conv3d(nf1, nf2, kernel_size=1, bias=self.use_bias),
nn.BatchNorm3d(nf2), nn.ReLU(True))
# 4 -> 8
nf3 = nf2 if not self.use_skip_dense else nf1 + nf2
nf4 = nf3 // 2
self.decode_dense3 = nn.Sequential(
nn.ConvTranspose3d(nf3,
nf4,
kernel_size=4,
stride=2,
padding=1,
bias=self.use_bias), nn.BatchNorm3d(nf4),
nn.ReLU(True))
# 8 -> 16
if self.use_skip_dense:
nf4 += nf0
nf5 = nf4 // 2
self.decode_dense4 = nn.Sequential(
nn.ConvTranspose3d(nf4,
nf5,
kernel_size=4,
stride=2,
padding=1,
bias=self.use_bias), nn.BatchNorm3d(nf5),
nn.ReLU(True))
self.final = nn.Sequential(
nn.Conv3d(nf5, nf_out, kernel_size=1, bias=self.use_bias),
nn.BatchNorm3d(nf_out), nn.ReLU(True))
# occ prediction
self.occpred = nn.Sequential(
nn.Conv3d(nf_out, 1, kernel_size=1, bias=self.use_bias))
def __init__(self, cfg):
super(DPN, self).__init__()
in_planes, out_planes = cfg['in_planes'], cfg['out_planes']
num_blocks, dense_depth = cfg['num_blocks'], cfg['dense_depth']
# self.in_planes = in_planes
# self.out_planes = out_planes
# self.num_blocks = num_blocks
# self.dense_depth = dense_depth
self.conv1 = nn.Conv3d(1,
24,
kernel_size=3,
stride=1,
padding=1,
bias=False)
self.bn1 = nn.BatchNorm3d(24)
self.last_planes = 24
self.layer1 = self._make_layer(in_planes[0],
out_planes[0],
num_blocks[0],
dense_depth[0],
stride=2) #stride=1)
self.layer2 = self._make_layer(in_planes[1],
out_planes[1],
num_blocks[1],
dense_depth[1],
stride=2)
self.layer3 = self._make_layer(in_planes[2],
out_planes[2],
num_blocks[2],
dense_depth[2],
stride=2)
self.layer4 = self._make_layer(in_planes[3],
out_planes[3],
num_blocks[3],
dense_depth[3],
stride=2)
self.last_planes = 216
self.layer5 = self._make_layer(128,
128,
num_blocks[2],
dense_depth[2],
stride=1)
self.last_planes = 224 + 3
self.layer6 = self._make_layer(224,
224,
num_blocks[1],
dense_depth[1],
stride=1)
self.linear = nn.Linear(out_planes[3] +
(num_blocks[3] + 1) * dense_depth[3], 2) #10)
self.last_planes = 120
self.path1 = nn.Sequential(
nn.ConvTranspose3d(self.last_planes,
self.last_planes,
kernel_size=2,
stride=2), nn.BatchNorm3d(self.last_planes),
nn.ReLU(inplace=True))
self.last_planes = 152
self.path2 = nn.Sequential(
nn.ConvTranspose3d(self.last_planes,
self.last_planes,
kernel_size=2,
stride=2), nn.BatchNorm3d(self.last_planes),
nn.ReLU(inplace=True))
self.drop = nn.Dropout3d(p=0.5, inplace=False)
self.output = nn.Sequential(
nn.Conv3d(248, 64, kernel_size=1),
nn.ReLU(),
#nn.Dropout3d(p = 0.3),
nn.Conv3d(64, 5 * len(config['anchors']), kernel_size=1))
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]
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 test_builder_to_backend_stress(
self,
use_cpu_only,
backend,
conv_dim,
padding,
DHWKdKhKw,
stride,
dilation,
has_bias,
groups,
test_symbolic,
test_output_shape,
):
D, H, W, Kd, Kh, Kw = DHWKdKhKw
N, C_in, C_out = 1, 1 * groups, 2 * groups
import torch
import torch.nn as nn
isDeconv1d = conv_dim == "conv1d"
isDeconv2d = conv_dim == "conv2d"
if isDeconv1d:
strides = [stride[0]]
dilations = [dilation[0]]
kernels = [Kh]
m = nn.ConvTranspose1d(
C_in,
C_out,
kernels,
stride=strides,
dilation=dilations,
bias=has_bias,
groups=groups,
padding=padding[0],
)
input_shape = [N, C_in, H]
paddings = [padding[0], padding[0]]
elif isDeconv2d:
strides = [stride[0], stride[1]]
dilations = [dilation[0], dilation[1]]
kernels = [Kh, Kw]
m = nn.ConvTranspose2d(
C_in,
C_out,
kernels,
stride=strides,
dilation=dilations,
bias=has_bias,
groups=groups,
padding=(padding[0], padding[1]),
)
input_shape = [N, C_in, H, W]
paddings = [padding[0], padding[0], padding[1], padding[1]]
else:
strides = [stride[0], stride[1], stride[2]]
dilations = [dilation[0], dilation[1], dilation[2]]
kernels = [Kd, Kh, Kw]
m = nn.ConvTranspose3d(
C_in,
C_out,
kernels,
stride=strides,
dilation=dilations,
bias=has_bias,
groups=groups,
padding=padding,
)
input_shape = [N, C_in, D, H, W]
paddings = [
padding[0],
padding[0],
padding[1],
padding[1],
padding[2],
padding[2],
]
wts = m.state_dict()
weight = wts["weight"].detach().numpy()
bias = wts["bias"].detach().numpy() if has_bias else None
# Reshape to CoreML format
# PyTorch weight format: C_in, C_out, H, W
# MIL weight format: C_out, C_in, H, W
if isDeconv1d:
weight = np.transpose(weight, [1, 0, 2])
elif isDeconv2d:
weight = np.transpose(weight, [1, 0, 2, 3])
else:
weight = np.transpose(weight, [1, 0, 2, 3, 4])
input = torch.randn(*input_shape)
output = m(input)
output = output.detach().numpy()
input = input.detach().numpy()
output_shape = list(output.shape)
if test_symbolic:
# For symbolic input test
# Make Batch Size and input channel as symbolic
symbolic_batch_size = get_new_symbol()
input_shape[0] = symbolic_batch_size
output_shape[0] = symbolic_batch_size
expected_output_types = tuple(output_shape[:]) + (types.fp32, )
expected_outputs = [output]
input_placeholders = {"x": mb.placeholder(shape=input_shape)}
input_values = {"x": input}
def build(x):
arguments = {
"x": x,
"weight": weight,
"pad": paddings,
"pad_type": "custom",
"strides": strides,
"dilations": dilations,
"groups": groups,
}
if has_bias:
arguments["bias"] = bias
if test_output_shape:
arguments["output_shape"] = output.shape[2:]
return mb.conv_transpose(**arguments)
run_compare_builder(
build,
input_placeholders,
input_values,
expected_output_types,
expected_outputs,
use_cpu_only=use_cpu_only,
frontend_only=False,
backend=backend,
)
