Beispiel #1
0
 def __init__(self, in_dim):
     super(NonLocal, self).__init__()
     self.chanel_in = in_dim
     self.theta = sn(nn.Conv2d(in_dim, in_dim // 8, 1, 1, 0))
     self.phi = sn(nn.Conv2d(in_dim, in_dim // 8, 1, 1, 0))
     self.g = sn(nn.Conv2d(in_dim, in_dim // 2, 1, 1, 0))
     self.out_conv = sn(nn.Conv2d(in_dim // 2, in_dim, 1, 1, 0))
     self.gamma = 1.0
     self.softmax = nn.Softmax(dim=-1)
     self.max_pool = nn.MaxPool2d(2)
Beispiel #2
0
 def __init__(self, fin, fout, cond_dim):
     super(ResBlock, self).__init__()
     self.learned_shortcut = (fin != fout)
     self.conv_0 = sn(nn.Conv2d(fin, fout, 3, 1, 1))
     self.conv_1 = sn(nn.Conv2d(fout, fout, 3, 1, 1))
     if self.learned_shortcut:
         self.conv_s = sn(nn.Conv2d(fin, fout, 1, 1, 0))
     self.relu_0 = nn.ReLU(True)
     self.relu_1 = nn.ReLU(True)
     self.bn_0 = CondBN(fin, cond_dim)
     self.bn_1 = CondBN(fout, cond_dim)
Beispiel #3
0
 def __init__(self, fin, fout, downsample):
     super(ResBlock, self).__init__()
     self.learned_shortcut = (fin != fout) or downsample
     self.downsample = downsample
     self.conv_0 = sn(nn.Conv2d(fin, fout, 3, 1, 1))
     self.conv_1 = sn(nn.Conv2d(fout, fout, 3, 1, 1))
     if self.learned_shortcut:
         self.conv_s = sn(nn.Conv2d(fin, fout, 1, 1, 0))
     self.relu_0 = nn.ReLU(False)
     self.relu_1 = nn.ReLU(False)
     if downsample:
         self.down_0 = nn.AvgPool2d(2)
         self.down_s = nn.AvgPool2d(2)
Beispiel #4
0
    def __init__(self):
        super(Generator, self).__init__()
        self.embedding_size = 80
        self.z_dim_chunk = 20
        self.num_chunk = 4
        self.cond_dim = self.z_dim_chunk + self.embedding_size

        self.fc = sn(nn.Linear(self.z_dim_chunk, 4 * 4 * 4 * nf))
        self.embedding = sn(nn.Linear(1000, self.embedding_size))

        self.resblock_0 = ResBlock(4 * nf, 4 * nf, cond_dim=self.cond_dim)
        self.resblock_1 = ResBlock(4 * nf, 2 * nf, cond_dim=self.cond_dim)
        self.non_local = NonLocal(2 * nf)
        self.resblock_2 = ResBlock(2 * nf, 1 * nf, cond_dim=self.cond_dim)

        self.bn_img = nn.BatchNorm2d(nf)
        self.relu_img = nn.ReLU(True)
        self.conv_img = sn(nn.Conv2d(nf, 3, 3, 1, 1))

        self.apply(init_func)
Beispiel #5
0
    def __init__(self):
        super(Discriminator, self).__init__()
        # prev conv
        self.pre_conv = nn.Sequential(sn(nn.Conv2d(3, 1 * nf, 3, 1, 1)),
                                      nn.ReLU(),
                                      sn(nn.Conv2d(1 * nf, 1 * nf, 3, 1, 1)),
                                      nn.AvgPool2d(2))
        self.pre_skip = nn.Sequential(
            nn.AvgPool2d(2),
            sn(nn.Conv2d(3, 1 * nf, 1, 1, 0)),
        )
        # body
        self.body = nn.Sequential(ResBlock(1 * nf, 2 * nf, downsample=True),
                                  NonLocal(2 * nf),
                                  ResBlock(2 * nf, 4 * nf, downsample=True),
                                  ResBlock(4 * nf, 4 * nf, downsample=False),
                                  nn.ReLU(True), nn.AdaptiveAvgPool2d(1))
        self.linear = sn(nn.Linear(4 * nf, 1))
        # embedding
        self.embed = nn.Embedding(1000, 4 * nf)
        self.embed = sn(self.embed)

        self.apply(init_func)
Beispiel #6
0
 def __init__(self, nc, cond_dim):
     super().__init__()
     self.mlp_gamma = sn(nn.Linear(cond_dim, nc))
     self.mlp_beta = sn(nn.Linear(cond_dim, nc))
     self.batch_norm = nn.BatchNorm2d(nc, momentum=0.001, affine=False)