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 )