def __init__(self, in_features): super(ResidualBlock, self).__init__() self.block = nn.Sequential(nn.ReflectionPad2d(1), nn.Conv(in_features, in_features, 3), nn.InstanceNorm2d(in_features, affine=None), nn.ReLU(), nn.ReflectionPad2d(1), nn.Conv(in_features, in_features, 3), nn.InstanceNorm2d(in_features, affine=None))
def test_batchnorm(self): # *************************************************************** # Test BatchNorm Layer # *************************************************************** arr = np.random.randn(16,10,224,224) check_equal_with_istrain(arr, jnn.BatchNorm(10, is_train=True), tnn.BatchNorm2d(10)) class Model(tnn.Module): def __init__(self): super(Model, self).__init__() self.layer = tnn.BatchNorm2d(10) def forward(self, x): return self.layer(x) model = Model() model.eval() check_equal_with_istrain(arr, jnn.BatchNorm(10, is_train=False), model, False) # *************************************************************** # Test InstanceNorm2d Layer # *************************************************************** arr = np.random.randn(16,10,224,224) check_equal_without_istrain(arr, jnn.InstanceNorm2d(10, is_train=True), tnn.InstanceNorm2d(10)) class Model(tnn.Module): def __init__(self): super(Model, self).__init__() self.layer = tnn.InstanceNorm2d(10) def forward(self, x): return self.layer(x) model = Model() model.eval() check_equal_without_istrain(arr, jnn.InstanceNorm2d(10, is_train=False), model) # *************************************************************** # Test BatchNorm1d Layer # *************************************************************** arr = np.random.randn(16,10) check_equal_with_istrain(arr, jnn.BatchNorm1d(10, is_train=True), tnn.BatchNorm1d(10), 1e-3) class Model(tnn.Module): def __init__(self): super(Model, self).__init__() self.layer = tnn.BatchNorm1d(10) def forward(self, x): return self.layer(x) model = Model() model.eval() check_equal_with_istrain(arr, jnn.BatchNorm1d(10, is_train=False), model, False)
def __init__(self, in_channel, style_dim): super(AdaptiveInstanceNorm, self).__init__() self.norm = nn.InstanceNorm2d(in_channel) self.linear = nn.Linear(style_dim, in_channel * 2) self.linear.bias.data[:in_channel] = 1 self.linear.bias.data[in_channel:] = 0
def discriminator_block(in_filters, out_filters, stride, normalize): 'Returns layers of each discriminator block' layers = [nn.Conv(in_filters, out_filters, 3, stride=stride, padding=1)] if normalize: layers.append(nn.InstanceNorm2d(out_filters)) layers.append(nn.Leaky_relu(scale=0.2)) return layers
def discriminator_block(in_filters, out_filters, normalization=True): 'Returns downsampling layers of each discriminator block' layers = [nn.Conv(in_filters, out_filters, 4, stride=2, padding=1)] if normalization: layers.append(nn.InstanceNorm2d(out_filters, affine=None)) layers.append(nn.LeakyReLU(scale=0.2)) return layers
def __init__(self, num_features, eps=1e-5, momentum=0.1): super(AdaptiveInstanceNorm2d, self).__init__() self.num_features = num_features self.norm = nn.InstanceNorm2d(num_features, affine=False) # weight and bias are dynamically assigned self.weight = None self.bias = None
def block(in_features, out_features, normalization=True): 'Classifier block' layers = [ nn.Conv(in_features, out_features, 3, stride=2, padding=1), nn.LeakyReLU(scale=0.2) ] if normalization: layers.append(nn.InstanceNorm2d(out_features, affine=None)) return layers
def __init__(self, in_size, out_size, normalize=True, dropout=0.0): super(UNetDown, self).__init__() layers = [ nn.Conv(in_size, out_size, 4, stride=2, padding=1, bias=False) ] if normalize: layers.append(nn.InstanceNorm2d(out_size, affine=None)) layers.append(nn.LeakyReLU(scale=0.2)) if dropout: layers.append(nn.Dropout(dropout)) self.model = nn.Sequential(*layers)
def __init__(self, img_shape=(3, 128, 128), res_blocks=9, c_dim=5): super(GeneratorResNet, self).__init__() (channels, img_size, _) = img_shape model = [ nn.Conv((channels + c_dim), 64, 7, stride=1, padding=3, bias=False), nn.InstanceNorm2d(64, affine=None), nn.ReLU() ] curr_dim = 64 for _ in range(2): model += [ nn.Conv(curr_dim, (curr_dim * 2), 4, stride=2, padding=1, bias=False), nn.InstanceNorm2d((curr_dim * 2), affine=None), nn.ReLU() ] curr_dim *= 2 for _ in range(res_blocks): model += [ResidualBlock(curr_dim)] for _ in range(2): model += [ nn.ConvTranspose(curr_dim, (curr_dim // 2), 4, stride=2, padding=1, bias=False), nn.InstanceNorm2d((curr_dim // 2), affine=None), nn.ReLU() ] curr_dim = (curr_dim // 2) model += [ nn.Conv(curr_dim, channels, 7, stride=1, padding=3), nn.Tanh() ] self.model = nn.Sequential(*model) for m in self.model: weights_init_normal(m)
def __init__(self, in_features): super(ResidualBlock, self).__init__() conv_block = [ nn.Conv(in_features, in_features, 3, stride=1, padding=1, bias=False), nn.InstanceNorm2d(in_features, affine=None), nn.ReLU(), nn.Conv(in_features, in_features, 3, stride=1, padding=1, bias=False), nn.InstanceNorm2d(in_features, affine=None) ] self.conv_block = nn.Sequential(*conv_block) for m in self.conv_block: weights_init_normal(m)
def __init__(self, input_shape, num_residual_blocks): super(GeneratorResNet, self).__init__() channels = input_shape[0] out_features = 64 model = [ nn.ReflectionPad2d(channels), nn.Conv(channels, out_features, 7), nn.InstanceNorm2d(out_features, affine=None), nn.ReLU() ] in_features = out_features for _ in range(2): out_features *= 2 model += [ nn.Conv(in_features, out_features, 3, stride=2, padding=1), nn.InstanceNorm2d(out_features, affine=None), nn.ReLU() ] in_features = out_features for _ in range(num_residual_blocks): model += [ResidualBlock(out_features)] for _ in range(2): out_features //= 2 model += [ nn.Upsample(scale_factor=2), nn.Conv(in_features, out_features, 3, stride=1, padding=1), nn.InstanceNorm2d(out_features, affine=None), nn.ReLU() ] in_features = out_features model += [ nn.ReflectionPad2d(channels), nn.Conv(out_features, channels, 7), nn.Tanh() ] self.model = nn.Sequential(*model) for m in self.modules(): weights_init_normal(m)
def __init__(self, input_dim, output_dim, kernel_size, stride, padding=0, norm='none', activation='relu', pad_type='zero'): super(ConvBlock, self).__init__() self.use_bias = True # initialize padding if pad_type == 'reflect': self.pad = nn.ReflectionPad2d(padding) elif pad_type == 'replicate': self.pad = nn.ReplicationPad2d(padding) elif pad_type == 'zero': self.pad = nn.ZeroPad2d(padding) else: assert 0, "Unsupported padding type: {}".format(pad_type) # initialize normalization norm_dim = output_dim if norm == 'bn': self.norm = nn.BatchNorm(norm_dim) elif norm == 'in': self.norm = nn.InstanceNorm2d(norm_dim) elif norm == 'adain': self.norm = AdaptiveInstanceNorm2d(norm_dim) elif norm == 'none': self.norm = None else: assert 0, "Unsupported normalization: {}".format(norm) # initialize activation if activation == 'relu': self.activation = nn.ReLU() elif activation == 'tanh': self.activation = nn.Tanh() elif activation == 'none': self.activation = None else: assert 0, "Unsupported activation: {}".format(activation) self.conv = nn.Conv(input_dim, output_dim, kernel_size, stride, bias=self.use_bias)
def __init__(self, in_size, out_size, dropout=0.0): super(UNetUp, self).__init__() layers = [ nn.ConvTranspose(in_size, out_size, 4, stride=2, padding=1, bias=False), nn.InstanceNorm2d(out_size, affine=None), nn.ReLU() ] if dropout: layers.append(nn.Dropout(dropout)) self.model = nn.Sequential(*layers)
def __init__(self, channels, dlatent_size, use_wscale, use_noise, use_pixel_norm, use_instance_norm, use_styles, activation_layer): super().__init__() layers = [] if use_noise: layers.append(('noise', NoiseLayer(channels))) layers.append(('activation', activation_layer)) if use_pixel_norm: layers.append(('pixel_norm', PixelNormLayer())) if use_instance_norm: layers.append( ('instance_norm', nn.InstanceNorm2d(channels, affine=False))) self.top_epi = nn.Sequential(OrderedDict(layers)) if use_styles: self.style_mod = StyleMod(dlatent_size, channels, use_wscale=use_wscale) else: self.style_mod = None