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)
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)
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)
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)
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)
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)
