def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias):
conv_block = []
p = 0
if padding_type == 'reflect':
conv_block += [nn.ReflectionPad3d(1)]
elif padding_type == 'replicate':
conv_block += [nn.ReplicationPad3d(1)]
elif padding_type == 'zero':
p = 1
else:
raise NotImplementedError('padding [%s] is not implemented' % padding_type)
conv_block += [nn.Conv3d(dim, dim, kernel_size=3, padding=p, bias=use_bias),
norm_layer(dim),
nn.ReLU(True)]
if use_dropout:
conv_block += [nn.Dropout(0.5)]
p = 0
if padding_type == 'reflect':
conv_block += [nn.ReflectionPad3d(1)]
elif padding_type == 'replicate':
conv_block += [nn.ReplicationPad3d(1)]
elif padding_type == 'zero':
p = 1
else:
raise NotImplementedError('padding [%s] is not implemented' % padding_type)
conv_block += [nn.Conv3d(dim, dim, kernel_size=3, padding=p, bias=use_bias),
norm_layer(dim)]
return nn.Sequential(*conv_block)
def build_conv_block(self, dim, padding_type, norm_layer, use_dropout,
use_bias):
"""Construct a convolutional block.
Parameters:
dim (int) -- the number of channels in the conv layer.
padding_type (str) -- the name of padding layer: reflect | replicate | zero
norm_layer -- normalization layer
use_dropout (bool) -- if use dropout layers.
use_bias (bool) -- if the conv layer uses bias or not
Returns a conv block (with a conv layer, a normalization layer, and a non-linearity layer (ReLU))
"""
conv_block = []
p = 0
if padding_type == 'reflect':
conv_block += [nn.ReflectionPad3d(1)]
elif padding_type == 'replicate':
conv_block += [nn.ReplicationPad3d(1)]
elif padding_type == 'zero':
p = 1
else:
raise NotImplementedError('padding [%s] is not implemented' %
padding_type)
conv_block += [
nn.Conv3d(dim, dim, kernel_size=3, padding=p, bias=use_bias),
norm_layer(dim),
nn.ReLU(True)
]
if use_dropout:
conv_block += [nn.Dropout(0.5)]
p = 0
if padding_type == 'reflect':
conv_block += [nn.ReflectionPad3d(1)]
elif padding_type == 'replicate':
conv_block += [nn.ReplicationPad3d(1)]
elif padding_type == 'zero':
p = 1
else:
raise NotImplementedError('padding [%s] is not implemented' %
padding_type)
conv_block += [
nn.Conv3d(dim, dim, kernel_size=3, padding=p, bias=use_bias),
norm_layer(dim)
]
return nn.Sequential(*conv_block)
def __init__(self,
in_channels,
out_channels,
kernel_size,
sigma,
stride=1,
padding=0,
bias=None,
groups=3):
super(GaussianLayer3D, self).__init__()
if isinstance(kernel_size, tuple):
origin_kernel = np.zeros(kernel_size)
center = kernel_size[0] // 2
origin_kernel[center, center, center] = 1
else:
origin_kernel = np.zeros((kernel_size, kernel_size, kernel_size))
center = kernel_size // 2
origin_kernel[center, center, center] = 1
kernel = ndimage.gaussian_filter(origin_kernel, sigma=sigma)
kernel = torch.Tensor(kernel).unsqueeze(0).unsqueeze(0)
kernel = kernel.repeat(out_channels, in_channels, 1, 1, 1)
self.register_buffer('gaussian_kernel', kernel)
self.pad = nn.ReflectionPad3d(center)
self.gaussian = nn.Conv3d(in_channels,
out_channels,
kernel_size,
stride=stride,
padding=padding,
bias=bias,
groups=groups)
self.weights_init()
def conv3d_mod(in_f, out_f, kernel_size, stride=1, bias=True, pad='zero', downsample_mode='stride'):
downsampler = None
if stride != 1 and downsample_mode != 'stride':
if downsample_mode == 'avg':
downsampler = nn.AvgPool3d(stride, stride)
elif downsample_mode == 'max':
downsampler = nn.MaxPool3d(stride, stride)
elif downsample_mode in ['lanczos2', 'lanczos3']:
downsampler = Downsampler(n_planes=out_f, factor=stride, kernel_type=downsample_mode, phase=0.5,
preserve_size=True)
else:
assert False
stride = 1
padder = None
to_pad = int((kernel_size - 1) / 2)
if pad == 'reflection':
padder = nn.ReflectionPad3d(to_pad)
to_pad = 0
convolver = nn.Conv3d(in_f, out_f, kernel_size, stride, padding=to_pad, bias=bias)
layers = filter(lambda x: x is not None, [padder, convolver, downsampler])
return nn.Sequential(*layers)
def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm3d, use_dropout=False, n_blocks=6, padding_type='reflect'):
assert(n_blocks >= 0)
super(ResnetGenerator, self).__init__()
self.input_nc = input_nc
self.output_nc = output_nc
self.ngf = ngf
if type(norm_layer) == functools.partial:
use_bias = norm_layer.func == nn.InstanceNorm3d
else:
use_bias = norm_layer == nn.InstanceNorm3d
model = [nn.ReflectionPad3d(3),
nn.Conv3d(input_nc, ngf, kernel_size=7, padding=0,
bias=use_bias),
norm_layer(ngf),
nn.ReLU(True)]
n_downsampling = 2
for i in range(n_downsampling):
mult = 2**i
model += [nn.Conv3d(ngf * mult, ngf * mult * 2, kernel_size=3,
stride=2, padding=1, bias=use_bias),
norm_layer(ngf * mult * 2),
nn.ReLU(True)]
mult = 2**n_downsampling
for i in range(n_blocks):
model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]
for i in range(n_downsampling):
mult = 2**(n_downsampling - i)
model += [nn.ConvTranspose3d(ngf * mult, int(ngf * mult / 2),
kernel_size=3, stride=2,
padding=1, output_padding=1,
bias=use_bias),
norm_layer(int(ngf * mult / 2)),
nn.ReLU(True)]
model += [nn.ReflectionPad3d(3)]
model += [nn.Conv3d(ngf, output_nc, kernel_size=7, padding=0)]
model += [nn.Tanh()]
self.model = nn.Sequential(*model)
def pad(pad_type, padding):
# helper selecting padding layer
# if padding is 'zero', do by conv layers
pad_type = pad_type.lower()
if padding == 0:
return None
if pad_type == 'reflect':
layer = nn.ReflectionPad3d(padding)
elif pad_type == 'replicate':
layer = nn.ReplicationPad3d(padding)
else:
raise NotImplementedError(
'padding layer [{:s}] is not implemented'.format(pad_type))
return layer
def __init__(self):
super(NNPaddingModule, self).__init__()
self.input1d = torch.randn(1, 4, 50)
self.module1d = nn.ModuleList([
nn.ReflectionPad1d(2),
nn.ReplicationPad1d(2),
nn.ConstantPad1d(2, 3.5),
])
self.input2d = torch.randn(1, 4, 30, 10)
self.module2d = nn.ModuleList([
nn.ReflectionPad2d(2),
nn.ReplicationPad2d(2),
nn.ZeroPad2d(2),
nn.ConstantPad2d(2, 3.5),
])
self.input3d = torch.randn(1, 4, 10, 4, 4)
self.module3d = nn.ModuleList([
nn.ReflectionPad3d(1),
nn.ReplicationPad3d(3),
nn.ConstantPad3d(3, 3.5),
])
def __init__(self,
input_nc,
output_nc,
ngf=64,
norm_layer=nn.BatchNorm3d,
use_dropout=False,
n_blocks=6,
padding_type='reflect'):
"""Construct a Resnet-based generator
Parameters:
input_nc (int) -- the number of channels in input images
output_nc (int) -- the number of channels in output images
ngf (int) -- the number of filters in the last conv layer
norm_layer -- normalization layer
use_dropout (bool) -- if use dropout layers
n_blocks (int) -- the number of ResNet blocks
padding_type (str) -- the name of padding layer in conv layers: reflect | replicate | zero
"""
assert (n_blocks >= 0)
super(ResnetGenerator, self).__init__()
if type(norm_layer) == functools.partial:
use_bias = norm_layer.func == nn.InstanceNorm3d
else:
use_bias = norm_layer == nn.InstanceNorm3d
model = [
nn.ReflectionPad3d(3),
nn.Conv3d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),
norm_layer(ngf),
nn.ReLU(True)
]
n_downsampling = 2
for i in range(n_downsampling): # add downsampling layers
mult = 2**i
model += [
nn.Conv3d(ngf * mult,
ngf * mult * 2,
kernel_size=3,
stride=2,
padding=1,
bias=use_bias),
norm_layer(ngf * mult * 2),
nn.ReLU(True)
]
mult = 2**n_downsampling
for i in range(n_blocks): # add ResNet blocks
model += [
ResnetBlock(ngf * mult,
padding_type=padding_type,
norm_layer=norm_layer,
use_dropout=use_dropout,
use_bias=use_bias)
]
for i in range(n_downsampling): # add upsampling layers
mult = 2**(n_downsampling - i)
model += [
nn.ConvTranspose3d(ngf * mult,
int(ngf * mult / 2),
kernel_size=3,
stride=2,
padding=1,
output_padding=1,
bias=use_bias),
norm_layer(int(ngf * mult / 2)),
nn.ReLU(True)
]
model += [nn.ReflectionPad3d(3)]
model += [nn.Conv3d(ngf, output_nc, kernel_size=7, padding=0)]
model += [nn.Tanh()]
self.model = nn.Sequential(*model)
def __init__(self,
input_dim,
output_dim,
kernel_size,
stride,
padding=0,
norm='in',
activation='relu',
pad_type='zero'):
super(Conv3dBlock, self).__init__()
self.use_bias = True
# initialize padding
#############Pad Type no function for 3D, reflectionpad and zeropad #########
if pad_type == 'reflect':
self.pad = nn.ReflectionPad3d(padding)
elif pad_type == 'replicate':
self.pad = nn.ReplicationPad3d(padding)
elif pad_type == 'zero':
self.pad = nn.ZeroPad3d(padding)
else:
assert 0, "Unsupported padding type: {}".format(pad_type)
# initialize normalization
norm_dim = output_dim
if norm == 'bn':
self.norm = nn.BatchNorm3d(norm_dim)
elif norm == 'in':
#self.norm = nn.InstanceNorm2d(norm_dim, track_running_stats=True)
self.norm = nn.InstanceNorm3d(norm_dim)
elif norm == 'ln':
self.norm = LayerNorm(norm_dim)
elif norm == 'adain':
self.norm = AdaptiveInstanceNorm3d(norm_dim)
elif norm == 'none' or norm == 'sn':
self.norm = None
else:
assert 0, "Unsupported normalization: {}".format(norm)
# initialize activation
if activation == 'relu':
self.activation = nn.ReLU(inplace=True)
elif activation == 'lrelu':
self.activation = nn.LeakyReLU(0.2, inplace=True)
elif activation == 'prelu':
self.activation = nn.PReLU()
elif activation == 'selu':
self.activation = nn.SELU(inplace=True)
elif activation == 'tanh':
self.activation = nn.Tanh()
elif activation == 'none':
self.activation = None
else:
assert 0, "Unsupported activation: {}".format(activation)
# initialize convolution
if norm == 'sn':
self.conv = SpectralNorm(
nn.Conv3d(input_dim,
output_dim,
kernel_size,
stride,
bias=self.use_bias))
else:
self.conv = nn.Conv3d(input_dim,
output_dim,
kernel_size,
stride,
bias=self.use_bias)
