| | import torch |
| | import torch.nn as nn |
| | from torch.nn.utils import spectral_norm |
| | import torch.nn.functional as F |
| |
|
| | import random |
| |
|
| | seq = nn.Sequential |
| |
|
| | def weights_init(m): |
| | classname = m.__class__.__name__ |
| | if classname.find('Conv') != -1: |
| | try: |
| | m.weight.data.normal_(0.0, 0.02) |
| | except: |
| | pass |
| | elif classname.find('BatchNorm') != -1: |
| | m.weight.data.normal_(1.0, 0.02) |
| | m.bias.data.fill_(0) |
| |
|
| | def conv2d(*args, **kwargs): |
| | return spectral_norm(nn.Conv2d(*args, **kwargs)) |
| |
|
| | def convTranspose2d(*args, **kwargs): |
| | return spectral_norm(nn.ConvTranspose2d(*args, **kwargs)) |
| |
|
| | def batchNorm2d(*args, **kwargs): |
| | return nn.BatchNorm2d(*args, **kwargs) |
| |
|
| | def linear(*args, **kwargs): |
| | return spectral_norm(nn.Linear(*args, **kwargs)) |
| |
|
| | class PixelNorm(nn.Module): |
| | def forward(self, input): |
| | return input * torch.rsqrt(torch.mean(input ** 2, dim=1, keepdim=True) + 1e-8) |
| |
|
| | class Reshape(nn.Module): |
| | def __init__(self, shape): |
| | super().__init__() |
| | self.target_shape = shape |
| |
|
| | def forward(self, feat): |
| | batch = feat.shape[0] |
| | return feat.view(batch, *self.target_shape) |
| |
|
| |
|
| | class GLU(nn.Module): |
| | def forward(self, x): |
| | nc = x.size(1) |
| | assert nc % 2 == 0, 'channels dont divide 2!' |
| | nc = int(nc/2) |
| | return x[:, :nc] * torch.sigmoid(x[:, nc:]) |
| |
|
| |
|
| | class NoiseInjection(nn.Module): |
| | def __init__(self): |
| | super().__init__() |
| |
|
| | self.weight = nn.Parameter(torch.zeros(1), requires_grad=True) |
| |
|
| | def forward(self, feat, noise=None): |
| | if noise is None: |
| | batch, _, height, width = feat.shape |
| | noise = torch.randn(batch, 1, height, width).to(feat.device) |
| |
|
| | return feat + self.weight * noise |
| |
|
| |
|
| | class Swish(nn.Module): |
| | def forward(self, feat): |
| | return feat * torch.sigmoid(feat) |
| |
|
| |
|
| | class SEBlock(nn.Module): |
| | def __init__(self, ch_in, ch_out): |
| | super().__init__() |
| |
|
| | self.main = nn.Sequential( nn.AdaptiveAvgPool2d(4), |
| | conv2d(ch_in, ch_out, 4, 1, 0, bias=False), Swish(), |
| | conv2d(ch_out, ch_out, 1, 1, 0, bias=False), nn.Sigmoid() ) |
| |
|
| | def forward(self, feat_small, feat_big): |
| | return feat_big * self.main(feat_small) |
| |
|
| |
|
| | class InitLayer(nn.Module): |
| | def __init__(self, nz, channel): |
| | super().__init__() |
| |
|
| | self.init = nn.Sequential( |
| | convTranspose2d(nz, channel*2, 4, 1, 0, bias=False), |
| | batchNorm2d(channel*2), GLU() ) |
| |
|
| | def forward(self, noise): |
| | noise = noise.view(noise.shape[0], -1, 1, 1) |
| | return self.init(noise) |
| |
|
| |
|
| | def UpBlock(in_planes, out_planes): |
| | block = nn.Sequential( |
| | nn.Upsample(scale_factor=2, mode='nearest'), |
| | conv2d(in_planes, out_planes*2, 3, 1, 1, bias=False), |
| | |
| | batchNorm2d(out_planes*2), GLU()) |
| | return block |
| |
|
| |
|
| | def UpBlockComp(in_planes, out_planes): |
| | block = nn.Sequential( |
| | nn.Upsample(scale_factor=2, mode='nearest'), |
| | conv2d(in_planes, out_planes*2, 3, 1, 1, bias=False), |
| | |
| | NoiseInjection(), |
| | batchNorm2d(out_planes*2), GLU(), |
| | conv2d(out_planes, out_planes*2, 3, 1, 1, bias=False), |
| | NoiseInjection(), |
| | batchNorm2d(out_planes*2), GLU() |
| | ) |
| | return block |
| |
|
| |
|
| | class Generator(nn.Module): |
| | def __init__(self, ngf=64, nz=100, nc=3, im_size=1024): |
| | super(Generator, self).__init__() |
| |
|
| | nfc_multi = {4:16, 8:8, 16:4, 32:2, 64:2, 128:1, 256:0.5, 512:0.25, 1024:0.125} |
| | nfc = {} |
| | for k, v in nfc_multi.items(): |
| | nfc[k] = int(v*ngf) |
| |
|
| | self.im_size = im_size |
| |
|
| | self.init = InitLayer(nz, channel=nfc[4]) |
| | |
| | self.feat_8 = UpBlockComp(nfc[4], nfc[8]) |
| | self.feat_16 = UpBlock(nfc[8], nfc[16]) |
| | self.feat_32 = UpBlockComp(nfc[16], nfc[32]) |
| | self.feat_64 = UpBlock(nfc[32], nfc[64]) |
| | self.feat_128 = UpBlockComp(nfc[64], nfc[128]) |
| | self.feat_256 = UpBlock(nfc[128], nfc[256]) |
| |
|
| | self.se_64 = SEBlock(nfc[4], nfc[64]) |
| | self.se_128 = SEBlock(nfc[8], nfc[128]) |
| | self.se_256 = SEBlock(nfc[16], nfc[256]) |
| |
|
| | self.to_128 = conv2d(nfc[128], nc, 1, 1, 0, bias=False) |
| | self.to_big = conv2d(nfc[im_size], nc, 3, 1, 1, bias=False) |
| | |
| | if im_size > 256: |
| | self.feat_512 = UpBlockComp(nfc[256], nfc[512]) |
| | self.se_512 = SEBlock(nfc[32], nfc[512]) |
| | if im_size > 512: |
| | self.feat_1024 = UpBlock(nfc[512], nfc[1024]) |
| | |
| | def forward(self, input): |
| | |
| | feat_4 = self.init(input) |
| | feat_8 = self.feat_8(feat_4) |
| | feat_16 = self.feat_16(feat_8) |
| | feat_32 = self.feat_32(feat_16) |
| |
|
| | feat_64 = self.se_64( feat_4, self.feat_64(feat_32) ) |
| |
|
| | feat_128 = self.se_128( feat_8, self.feat_128(feat_64) ) |
| |
|
| | feat_256 = self.se_256( feat_16, self.feat_256(feat_128) ) |
| |
|
| | if self.im_size == 256: |
| | return [self.to_big(feat_256), self.to_128(feat_128)] |
| | |
| | feat_512 = self.se_512( feat_32, self.feat_512(feat_256) ) |
| | if self.im_size == 512: |
| | return [self.to_big(feat_512), self.to_128(feat_128)] |
| |
|
| | feat_1024 = self.feat_1024(feat_512) |
| |
|
| | im_128 = torch.tanh(self.to_128(feat_128)) |
| | im_1024 = torch.tanh(self.to_big(feat_1024)) |
| |
|
| | return [im_1024, im_128] |
| |
|
| |
|
| | class DownBlock(nn.Module): |
| | def __init__(self, in_planes, out_planes): |
| | super(DownBlock, self).__init__() |
| |
|
| | self.main = nn.Sequential( |
| | conv2d(in_planes, out_planes, 4, 2, 1, bias=False), |
| | batchNorm2d(out_planes), nn.LeakyReLU(0.2, inplace=True), |
| | ) |
| |
|
| | def forward(self, feat): |
| | return self.main(feat) |
| |
|
| |
|
| | class DownBlockComp(nn.Module): |
| | def __init__(self, in_planes, out_planes): |
| | super(DownBlockComp, self).__init__() |
| |
|
| | self.main = nn.Sequential( |
| | conv2d(in_planes, out_planes, 4, 2, 1, bias=False), |
| | batchNorm2d(out_planes), nn.LeakyReLU(0.2, inplace=True), |
| | conv2d(out_planes, out_planes, 3, 1, 1, bias=False), |
| | batchNorm2d(out_planes), nn.LeakyReLU(0.2) |
| | ) |
| |
|
| | self.direct = nn.Sequential( |
| | nn.AvgPool2d(2, 2), |
| | conv2d(in_planes, out_planes, 1, 1, 0, bias=False), |
| | batchNorm2d(out_planes), nn.LeakyReLU(0.2)) |
| |
|
| | def forward(self, feat): |
| | return (self.main(feat) + self.direct(feat)) / 2 |
| |
|
| |
|
| | class Discriminator(nn.Module): |
| | def __init__(self, ndf=64, nc=3, im_size=512): |
| | super(Discriminator, self).__init__() |
| | self.ndf = ndf |
| | self.im_size = im_size |
| |
|
| | nfc_multi = {4:16, 8:16, 16:8, 32:4, 64:2, 128:1, 256:0.5, 512:0.25, 1024:0.125} |
| | nfc = {} |
| | for k, v in nfc_multi.items(): |
| | nfc[k] = int(v*ndf) |
| |
|
| | if im_size == 1024: |
| | self.down_from_big = nn.Sequential( |
| | conv2d(nc, nfc[1024], 4, 2, 1, bias=False), |
| | nn.LeakyReLU(0.2, inplace=True), |
| | conv2d(nfc[1024], nfc[512], 4, 2, 1, bias=False), |
| | batchNorm2d(nfc[512]), |
| | nn.LeakyReLU(0.2, inplace=True)) |
| | elif im_size == 512: |
| | self.down_from_big = nn.Sequential( |
| | conv2d(nc, nfc[512], 4, 2, 1, bias=False), |
| | nn.LeakyReLU(0.2, inplace=True) ) |
| | elif im_size == 256: |
| | self.down_from_big = nn.Sequential( |
| | conv2d(nc, nfc[512], 3, 1, 1, bias=False), |
| | nn.LeakyReLU(0.2, inplace=True) ) |
| |
|
| | self.down_4 = DownBlockComp(nfc[512], nfc[256]) |
| | self.down_8 = DownBlockComp(nfc[256], nfc[128]) |
| | self.down_16 = DownBlockComp(nfc[128], nfc[64]) |
| | self.down_32 = DownBlockComp(nfc[64], nfc[32]) |
| | self.down_64 = DownBlockComp(nfc[32], nfc[16]) |
| |
|
| | self.rf_big = nn.Sequential( |
| | conv2d(nfc[16] , nfc[8], 1, 1, 0, bias=False), |
| | batchNorm2d(nfc[8]), nn.LeakyReLU(0.2, inplace=True), |
| | conv2d(nfc[8], 1, 4, 1, 0, bias=False)) |
| |
|
| | self.se_2_16 = SEBlock(nfc[512], nfc[64]) |
| | self.se_4_32 = SEBlock(nfc[256], nfc[32]) |
| | self.se_8_64 = SEBlock(nfc[128], nfc[16]) |
| | |
| | self.down_from_small = nn.Sequential( |
| | conv2d(nc, nfc[256], 4, 2, 1, bias=False), |
| | nn.LeakyReLU(0.2, inplace=True), |
| | DownBlock(nfc[256], nfc[128]), |
| | DownBlock(nfc[128], nfc[64]), |
| | DownBlock(nfc[64], nfc[32]), ) |
| |
|
| | self.rf_small = conv2d(nfc[32], 1, 4, 1, 0, bias=False) |
| |
|
| | self.decoder_big = SimpleDecoder(nfc[16], nc) |
| | self.decoder_part = SimpleDecoder(nfc[32], nc) |
| | self.decoder_small = SimpleDecoder(nfc[32], nc) |
| | |
| | def forward(self, imgs, label, part=None): |
| | if type(imgs) is not list: |
| | imgs = [F.interpolate(imgs, size=self.im_size), F.interpolate(imgs, size=128)] |
| |
|
| | feat_2 = self.down_from_big(imgs[0]) |
| | feat_4 = self.down_4(feat_2) |
| | feat_8 = self.down_8(feat_4) |
| | |
| | feat_16 = self.down_16(feat_8) |
| | feat_16 = self.se_2_16(feat_2, feat_16) |
| |
|
| | feat_32 = self.down_32(feat_16) |
| | feat_32 = self.se_4_32(feat_4, feat_32) |
| | |
| | feat_last = self.down_64(feat_32) |
| | feat_last = self.se_8_64(feat_8, feat_last) |
| |
|
| | |
| | |
| | rf_0 = self.rf_big(feat_last).view(-1) |
| |
|
| | feat_small = self.down_from_small(imgs[1]) |
| | |
| | rf_1 = self.rf_small(feat_small).view(-1) |
| |
|
| | if label=='real': |
| | rec_img_big = self.decoder_big(feat_last) |
| | rec_img_small = self.decoder_small(feat_small) |
| |
|
| | assert part is not None |
| | rec_img_part = None |
| | if part==0: |
| | rec_img_part = self.decoder_part(feat_32[:,:,:8,:8]) |
| | if part==1: |
| | rec_img_part = self.decoder_part(feat_32[:,:,:8,8:]) |
| | if part==2: |
| | rec_img_part = self.decoder_part(feat_32[:,:,8:,:8]) |
| | if part==3: |
| | rec_img_part = self.decoder_part(feat_32[:,:,8:,8:]) |
| |
|
| | return torch.cat([rf_0, rf_1]) , [rec_img_big, rec_img_small, rec_img_part] |
| |
|
| | return torch.cat([rf_0, rf_1]) |
| |
|
| |
|
| | class SimpleDecoder(nn.Module): |
| | """docstring for CAN_SimpleDecoder""" |
| | def __init__(self, nfc_in=64, nc=3): |
| | super(SimpleDecoder, self).__init__() |
| |
|
| | nfc_multi = {4:16, 8:8, 16:4, 32:2, 64:2, 128:1, 256:0.5, 512:0.25, 1024:0.125} |
| | nfc = {} |
| | for k, v in nfc_multi.items(): |
| | nfc[k] = int(v*32) |
| |
|
| | def upBlock(in_planes, out_planes): |
| | block = nn.Sequential( |
| | nn.Upsample(scale_factor=2, mode='nearest'), |
| | conv2d(in_planes, out_planes*2, 3, 1, 1, bias=False), |
| | batchNorm2d(out_planes*2), GLU()) |
| | return block |
| |
|
| | self.main = nn.Sequential( nn.AdaptiveAvgPool2d(8), |
| | upBlock(nfc_in, nfc[16]) , |
| | upBlock(nfc[16], nfc[32]), |
| | upBlock(nfc[32], nfc[64]), |
| | upBlock(nfc[64], nfc[128]), |
| | conv2d(nfc[128], nc, 3, 1, 1, bias=False), |
| | nn.Tanh() ) |
| |
|
| | def forward(self, input): |
| | |
| | return self.main(input) |
| |
|
| | from random import randint |
| | def random_crop(image, size): |
| | h, w = image.shape[2:] |
| | ch = randint(0, h-size-1) |
| | cw = randint(0, w-size-1) |
| | return image[:,:,ch:ch+size,cw:cw+size] |
| |
|
| | class TextureDiscriminator(nn.Module): |
| | def __init__(self, ndf=64, nc=3, im_size=512): |
| | super(TextureDiscriminator, self).__init__() |
| | self.ndf = ndf |
| | self.im_size = im_size |
| |
|
| | nfc_multi = {4:16, 8:8, 16:8, 32:4, 64:2, 128:1, 256:0.5, 512:0.25, 1024:0.125} |
| | nfc = {} |
| | for k, v in nfc_multi.items(): |
| | nfc[k] = int(v*ndf) |
| |
|
| | self.down_from_small = nn.Sequential( |
| | conv2d(nc, nfc[256], 4, 2, 1, bias=False), |
| | nn.LeakyReLU(0.2, inplace=True), |
| | DownBlock(nfc[256], nfc[128]), |
| | DownBlock(nfc[128], nfc[64]), |
| | DownBlock(nfc[64], nfc[32]), ) |
| | self.rf_small = nn.Sequential( |
| | conv2d(nfc[16], 1, 4, 1, 0, bias=False)) |
| |
|
| | self.decoder_small = SimpleDecoder(nfc[32], nc) |
| | |
| | def forward(self, img, label): |
| | img = random_crop(img, size=128) |
| |
|
| | feat_small = self.down_from_small(img) |
| | rf = self.rf_small(feat_small).view(-1) |
| | |
| | if label=='real': |
| | rec_img_small = self.decoder_small(feat_small) |
| |
|
| | return rf, rec_img_small, img |
| |
|
| | return rf |
