def __init__(self, block, layers): self.inplanes = 64 super(FPN, self).__init__() self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False) self.bn1 = nn.BatchNorm2d(64) self.relu = nn.ReLU(inplace=True) self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) self.layer1 = self._make_layer(block, 64, layers[0]) # conv2 C2 self.layer2 = self._make_layer(block, 128, layers[1], stride=2) # conv3 C3 self.layer3 = self._make_layer(block, 256, layers[2], stride=2) # conv4 C4 self.layer4 = self._make_layer(block, 512, layers[3], stride=2) # conv5 C5 self.conv6 = conv3x3(512 * block.expansion, 256, stride=2, padding=1) # P6 self.conv7 = conv3x3(256, 256, stride=2, padding=1) # P7 self.lateral_layer1 = conv1x1(512 * block.expansion, 256) self.lateral_layer2 = conv1x1(256 * block.expansion, 256) self.lateral_layer3 = conv1x1(128 * block.expansion, 256) self.corr_layer1 = conv3x3(256, 256, stride=1, padding=1) # P4 self.corr_layer2 = conv3x3(256, 256, stride=1, padding=1) # P3
def __init__(self, inplanes, planes, stride=1, downsample=None): super(BasicBlock, self).__init__() self.conv1 = conv3x3(inplanes, planes, stride) self.bn1 = nn.BatchNorm2d(planes) self.relu = nn.ReLU(inplace=True) self.conv2 = conv3x3(planes, planes) self.bn2 = nn.BatchNorm2d(planes) self.downsample = downsample self.stride = stride
def __init__(self, in_channel, out_channel, condition_dim, cross_replica=False, upsample=True): super().__init__() self.cbn1 = ConditionalNorm(in_channel, condition_dim, cross_replica) self.upsample = nn.Sequential() if upsample: self.upsample.add_module('upsample', nn.Upsample(scale_factor=2, mode='nearest')) self.conv3x3_1 = nn.utils.spectral_norm( conv3x3(in_channel, out_channel)).apply(init_weight) self.cbn2 = ConditionalNorm(out_channel, condition_dim, cross_replica) self.conv3x3_2 = nn.utils.spectral_norm( conv3x3(out_channel, out_channel)).apply(init_weight) self.conv1x1 = nn.utils.spectral_norm( conv1x1(in_channel, out_channel)).apply(init_weight)
def __init__(self, in_channel, out_channel, downsample=True): super().__init__() self.layer = nn.Sequential( nn.LeakyReLU(0.2), nn.utils.spectral_norm(conv3x3(in_channel, out_channel)).apply( init_weight), nn.LeakyReLU(0.2), nn.utils.spectral_norm(conv3x3(out_channel, out_channel)).apply( init_weight), ) self.shortcut = nn.Sequential( nn.utils.spectral_norm(conv1x1(in_channel, out_channel)).apply( init_weight), ) if downsample: self.layer.add_module('avgpool', nn.AvgPool2d(kernel_size=2, stride=2)) self.shortcut.add_module('avgpool', nn.AvgPool2d(kernel_size=2, stride=2))
def __init__(self, z_dim, conv_dim, num_classes): super(Generator, self).__init__() self.conv_dim = conv_dim self.linear = spectral_norm( linear(in_features=z_dim, out_features=conv_dim * 16 * 4 * 4)) self.res_1 = ResidualBlock_G(conv_dim * 16, conv_dim * 16, num_classes) self.res_2 = ResidualBlock_G(conv_dim * 16, conv_dim * 8, num_classes) self.res_3 = ResidualBlock_G(conv_dim * 8, conv_dim * 4, num_classes) self.attn = SelfAttn(conv_dim * 4) self.res_4 = ResidualBlock_G(conv_dim * 4, conv_dim * 2, num_classes) self.res_5 = ResidualBlock_G(conv_dim * 2, conv_dim, num_classes) self.bn = batch_norm(conv_dim, eps=1e-5, momentum=0.0001) self.lrelu = lrelu(inplace=True) self.conv3x3 = spectral_norm(conv3x3(conv_dim, 3)) self.tanh = tanh() self.apply(init_weights)
def __init__(self, n_feat, max_resolution, codes_dim, n_classes=0, use_attention=False, arch=None, cross_replica=False, rgb_bn=False, use_embedding=True): super().__init__() self.codes_dim = codes_dim self.use_embedding= use_embedding self.n_classes = n_classes # construct residual blocks n_layers = int(np.log2(max_resolution) - 2) self.residual_blocks = nn.ModuleList([]) last_block_dim = 0 if arch is None: first_block_factor = 2 ** (n_layers) else: first_block_factor = arch[0] self.fc = nn.Sequential( nn.utils.spectral_norm( nn.Linear(codes_dim, first_block_factor * n_feat * 4 * 4)).apply( init_weight) ) # print("first_block", first_block_factor) # print("n_layers ", n_layers) for i in range(n_layers): if arch is None: prev_factor = 2 ** (n_layers - i) curr_factor = 2 ** (n_layers - i - 1) else: prev_factor = arch[i] curr_factor = arch[i + 1] # print(f"block ({i}): {prev_factor}, {curr_factor}") block = ResBlock_G(prev_factor * n_feat, curr_factor * n_feat, codes_dim + codes_dim, cross_replica=cross_replica, upsample=True) # add current block to the model class self.residual_blocks.add_module(f'res_block_{i}', block) if i == n_layers - 1: last_block_dim = curr_factor if use_attention: self.attn = Attention(2 * n_feat) # print("last_layer ", last_block_dim) _to_rgb = [ # nn.BatchNorm2d(2 * n_feat).apply(init_weight), nn.LeakyReLU(), nn.utils.spectral_norm(conv3x3(last_block_dim * n_feat, 3)).apply( init_weight), nn.Tanh() ] bn_layer = SyncBN2d if cross_replica else nn.BatchNorm2d if rgb_bn: _to_rgb.insert(0, bn_layer(last_block_dim * n_feat)) self.to_rgb = nn.Sequential( *_to_rgb ) self.embedding = nn.Embedding(num_embeddings=n_classes, embedding_dim=self.codes_dim).apply(init_weight)