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)