def __init__(self, scale_factor, num_channels, num_features, growth_rate,
num_blocks, num_layers):
super(BlurRDN, self).__init__()
self.G0 = num_features
self.G = growth_rate
self.D = num_blocks
self.C = num_layers
# shallow feature extraction
self.sfe1 = nn.Conv2d(num_channels,
num_features,
kernel_size=3,
padding=3 // 2)
self.sfe2 = nn.Conv2d(num_features,
num_features,
kernel_size=3,
padding=3 // 2)
# residual dense blocks
self.rdbs = nn.ModuleList([RDB(self.G0, self.G, self.C)])
for _ in range(self.D - 1):
self.rdbs.append(RDB(self.G, self.G, self.C))
# global feature fusion
self.gff = nn.Sequential(
nn.Conv2d(self.G * self.D, self.G0, kernel_size=1),
nn.Conv2d(self.G0, self.G0, kernel_size=3, padding=3 // 2))
# up-sampling
assert 2 <= scale_factor <= 4
blurpool = antialiased_cnns.BlurPool(self.G0, filt_size=3, stride=1)
if scale_factor == 2 or scale_factor == 4:
self.upscale = []
for _ in range(scale_factor // 2):
self.upscale.extend([
nn.Conv2d(self.G0,
self.G0 * (2**2),
kernel_size=3,
padding=3 // 2),
nn.PixelShuffle(2)
])
self.upscale = nn.Sequential(*self.upscale, blurpool)
else:
self.upscale = nn.Sequential(
nn.Conv2d(self.G0,
self.G0 * (scale_factor**2),
kernel_size=3,
padding=3 // 2), nn.PixelShuffle(scale_factor),
blurpool)
self.output = nn.Conv2d(self.G0,
num_channels,
kernel_size=3,
padding=3 // 2)
def __init__(self,
imsize=64,
z_dim=100,
base_channels=32,
nc=3,
n_expert_components=5):
super(Enc, self).__init__()
self.imsize = imsize
self.n_expert_components = n_expert_components
self.z_dim = z_dim
self.nc = nc
self.encoder = nn.Sequential(
nn.Conv2d(in_channels=self.nc,
out_channels=32,
kernel_size=4,
stride=1,
padding=1), # B, 32, 32, 32
nn.LeakyReLU(True),
antialiased_cnns.BlurPool(32, filt_size=4, stride=2),
nn.Conv2d(32, 32, 4, 1, 1), # B, 32, 16, 16
nn.LeakyReLU(True),
antialiased_cnns.BlurPool(32, filt_size=4, stride=2),
nn.Conv2d(32, 64, 4, 1, 1), # B, 64, 8, 8
nn.LeakyReLU(True),
antialiased_cnns.BlurPool(64, filt_size=4, stride=2),
nn.Conv2d(64, 64, 4, 1, 1), # B, 64, 4, 4
nn.LeakyReLU(True),
antialiased_cnns.BlurPool(64, filt_size=4, stride=2),
nn.Conv2d(64, 256, 4, 1), # B, 256, 1, 1
nn.LeakyReLU(True),
View((-1, 256 * 1 * 1))) # B, 256)
self.muvar = nn.Linear(256, z_dim * 2)
self.BC_1 = nn.Linear(256, 128)
self.BC_2 = nn.Linear(128, 2 * self.n_expert_components)
# self.weight_init()
self.act = nn.LeakyReLU()
self.softplus = nn.Softplus()
self.sigmoid = nn.Sigmoid()
def create_stem(
in_chans=3, out_chs=32, kernel_size=3, stride=2, pool='',
act_layer=None, norm_layer=None, aa_layer=None):
stem = nn.Sequential()
if not isinstance(out_chs, (tuple, list)):
out_chs = [out_chs]
assert len(out_chs)
in_c = in_chans
for i, out_c in enumerate(out_chs):
conv_name = f'conv{i + 1}'
stem.add_module(conv_name, ConvBnAct(
in_c, out_c, kernel_size, stride=stride if i == 0 else 1,
act_layer=act_layer, norm_layer=norm_layer))
in_c = out_c
last_conv = conv_name
if pool:
if aa_layer is not None:
stem.add_module('pool', nn.MaxPool2d(kernel_size=3, stride=1, padding=1))
stem.add_module('aa', aa_layer(channels=in_c, stride=2))
else:
stem.add_module('pool', [nn.MaxPool2d(kernel_size=3, stride=1), antialiased_cnns.BlurPool(in_c, stride=2)])
'''nn.MaxPool2d(kernel_size=3, stride=2, padding=1)'''
return stem, dict(num_chs=in_c, reduction=stride, module='.'.join(['stem', last_conv]))
def __init__(self,
input_nc,
output_nc,
ngf=64,
norm_layer=nn.BatchNorm2d,
use_dropout=False,
n_blocks=6,
padding_type='reflect',
demodule=False,
anti_alias=False):
"""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.InstanceNorm2d
else:
use_bias = norm_layer == nn.InstanceNorm2d
if demodule:
model = [
nn.ReflectionPad2d(3),
ModulatedConv2d(input_nc, ngf, kernel_size=7),
norm_layer(ngf),
nn.ReLU(True)
]
else:
model = [
nn.ReflectionPad2d(3),
nn.Conv2d(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
if demodule:
model += [
ModulatedConv2d(ngf * mult,
ngf * mult * 2,
kernel_size=3,
downsample=True),
norm_layer(ngf * mult * 2),
nn.ReLU(True)
]
elif anti_alias:
model += [
nn.Conv2d(ngf * mult,
ngf * mult * 2,
kernel_size=3,
stride=1,
padding=1),
norm_layer(ngf * mult * 2),
nn.ReLU(inplace=True),
antialiased_cnns.BlurPool(ngf * mult * 2, stride=2),
]
else:
model += [
nn.Conv2d(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,
demodule=demodule)
]
for i in range(n_downsampling): # add upsampling layers
mult = 2**(n_downsampling - i)
# model += [nn.Upsample(scale_factor=2, mode='bilinear'),
# nn.ReflectionPad2d(1),
# nn.Conv2d(ngf * mult, int(ngf * mult / 2),
# kernel_size=3, stride=1, padding=0),norm_layer(int(ngf * mult / 2)),
# nn.ReLU(True)]
if demodule:
model += [
ModulatedConv2d(ngf * mult,
int(ngf * mult / 2),
kernel_size=3,
upsample=True),
norm_layer(int(ngf * mult / 2)),
nn.ReLU(True)
]
elif anti_alias:
model += [
nn.ConvTranspose2d(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(inplace=True),
antialiased_cnns.BlurPool(int(ngf * mult / 2), stride=1)
]
else:
model += [
nn.ConvTranspose2d(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.ReflectionPad2d(3)]
if demodule:
model += [ModulatedConv2d(ngf, output_nc, kernel_size=7)]
else:
model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
model += [nn.Tanh()]
self.model = nn.Sequential(*model)
def __init__(
self,
in_channels,
out_channels,
*,
bottleneck_channels,
stride=1,
num_groups=1,
norm="BN",
stride_in_1x1=False,
dilation=1,
avd=False,
avg_down=False,
radix=2,
bottleneck_width=64,
):
"""
Args:
norm (str or callable): a callable that takes the number of
channels and return a `nn.Module`, or a pre-defined string
(one of {"FrozenBN", "BN", "GN"}).
stride_in_1x1 (bool): when stride==2, whether to put stride in the
first 1x1 convolution or the bottleneck 3x3 convolution.
"""
super().__init__(in_channels, out_channels, stride)
self.avd = avd and (stride > 1)
self.avg_down = avg_down
self.radix = radix
cardinality = num_groups
group_width = int(bottleneck_channels *
(bottleneck_width / 64.)) * cardinality
filt_size = 7
if in_channels != out_channels:
if self.avg_down:
if stride == 1:
self.shortcut_avgpool = nn.AvgPool2d(
kernel_size=stride,
stride=stride,
ceil_mode=True,
count_include_pad=False)
else:
self.shortcut_avgpool = antialiased_cnns.BlurPool(
in_channels, stride=2, filt_size=filt_size)
print('shortcut_avgpool kernel size:', stride, ' stride:',
stride)
self.shortcut = Conv2d(
in_channels,
out_channels,
kernel_size=1,
stride=1,
bias=False,
norm=get_norm(norm, out_channels),
)
else:
self.shortcut = Conv2d(
in_channels,
out_channels,
kernel_size=1,
stride=stride,
bias=False,
norm=get_norm(norm, out_channels),
)
else:
self.shortcut = None
# The original MSRA ResNet models have stride in the first 1x1 conv
# The subsequent fb.torch.resnet and Caffe2 ResNe[X]t implementations have
# stride in the 3x3 conv
stride_1x1, stride_3x3 = (stride, 1) if stride_in_1x1 else (1, stride)
self.conv1 = Conv2d(
in_channels,
group_width,
kernel_size=1,
stride=stride_1x1,
bias=False,
norm=get_norm(norm, group_width),
)
if self.radix > 1:
from .splat import SplAtConv2d
self.conv2 = SplAtConv2d(
group_width,
group_width,
kernel_size=3,
stride=1 if self.avd else stride_3x3,
padding=dilation,
dilation=dilation,
groups=cardinality,
bias=False,
radix=self.radix,
norm=norm,
)
else:
self.conv2 = Conv2d(
group_width,
group_width,
kernel_size=3,
stride=1 if self.avd else stride_3x3,
padding=1 * dilation,
bias=False,
groups=num_groups,
dilation=dilation,
norm=get_norm(norm, group_width),
)
if self.avd:
#self.avd_layer = nn.AvgPool2d(3, stride, padding=1)
print('avd_layer stride:', stride)
self.avd_layer = nn.Sequential(nn.AvgPool2d(3, 1, padding=1), \
antialiased_cnns.BlurPool(group_width, stride=2, filt_size=filt_size))
self.conv3 = Conv2d(
group_width,
out_channels,
kernel_size=1,
bias=False,
norm=get_norm(norm, out_channels),
)
if self.radix > 1:
for layer in [self.conv1, self.conv3, self.shortcut]:
if layer is not None: # shortcut can be None
weight_init.c2_msra_fill(layer)
else:
for layer in [self.conv1, self.conv2, self.conv3, self.shortcut]:
if layer is not None: # shortcut can be None
weight_init.c2_msra_fill(layer)
def __init__(self, hidden):
super(Autoencoder, self).__init__()
self.encoder_conv = nn.Sequential(
nn.Conv2d(1, 16, 4, stride=1, padding=1),
nn.LeakyReLU(self.RELU_LEAK),
antialiased_cnns.BlurPool(16, stride=2),
nn.Conv2d(16, 32, 4, stride=1, padding=1),
nn.LeakyReLU(self.RELU_LEAK),
antialiased_cnns.BlurPool(32, stride=2),
nn.Conv2d(32, 64, 4, stride=1, padding=1),
nn.LeakyReLU(self.RELU_LEAK),
antialiased_cnns.BlurPool(64, stride=2),
nn.Conv2d(64, 128, 4, stride=1, padding=1),
nn.LeakyReLU(self.RELU_LEAK),
antialiased_cnns.BlurPool(128, stride=2),
nn.Conv2d(128, 256, 4, stride=1, padding=1),
nn.LeakyReLU(self.RELU_LEAK),
antialiased_cnns.BlurPool(256, stride=2),
)
self.encoder_linear = nn.Sequential(
nn.Linear(256 * 2 * 2, 256),
nn.LeakyReLU(self.RELU_LEAK),
)
self.encoder_mu = nn.Sequential(
nn.Linear(256, hidden), # No activation, let it go wherever
)
self.encoder_log_sigma = nn.Sequential(
nn.Linear(256, hidden), # Will be exponentiated later
)
self.decoder_linear = nn.Sequential(
nn.Linear(hidden, 256),
nn.LeakyReLU(self.RELU_LEAK),
nn.Linear(256, 256 * 2 * 2),
nn.LeakyReLU(self.RELU_LEAK),
)
self.decoder_layers = nn.ModuleList([
nn.Sequential(
nn.ConvTranspose2d(256,
128,
3,
stride=2,
padding=1,
output_padding=1),
nn.LeakyReLU(self.RELU_LEAK),
),
nn.Sequential(
nn.ConvTranspose2d(128,
64,
3,
stride=2,
padding=1,
output_padding=1),
nn.LeakyReLU(self.RELU_LEAK),
),
nn.Sequential(
nn.ConvTranspose2d(64,
32,
3,
stride=2,
padding=1,
output_padding=1),
nn.LeakyReLU(self.RELU_LEAK),
),
nn.Sequential(
nn.ConvTranspose2d(32,
16,
3,
stride=2,
padding=1,
output_padding=1),
nn.LeakyReLU(self.RELU_LEAK),
),
nn.Sequential(
nn.ConvTranspose2d(16,
8,
3,
stride=2,
padding=1,
output_padding=1),
nn.LeakyReLU(self.RELU_LEAK),
),
])
self.decoder_fc = nn.Sequential(
# Begin "fully connected layers"
nn.Conv2d(248, 512, 1),
nn.LeakyReLU(self.RELU_LEAK),
nn.Conv2d(512, 512, 1),
nn.LeakyReLU(self.RELU_LEAK),
nn.Conv2d(512, 1, 1),
nn.Sigmoid() # Can maybe clamp the output instead
)