| from torch import nn
|
| from torch.autograd import Variable
|
| import torch
|
| import torch.nn.functional as F
|
|
|
| import torchvision.models as models
|
|
|
| def add_gaussian_noise(ins, mean=0, stddev=0.1):
|
| noise = ins.data.new(ins.size()).normal_(mean, stddev)
|
| return ins + noise
|
|
|
| class FlattenLayer(nn.Module):
|
| def __init__(self):
|
| super(FlattenLayer, self).__init__()
|
|
|
| def forward(self, x):
|
| return x.view(x.size(0), -1)
|
|
|
|
|
| class UnflattenLayer(nn.Module):
|
| def __init__(self, width):
|
| super(UnflattenLayer, self).__init__()
|
| self.width = width
|
|
|
| def forward(self, x):
|
| return x.view(x.size(0), -1, self.width, self.width)
|
|
|
| class VAE_Encoder(nn.Module):
|
| '''
|
| VAE_Encoder: Encode image into std and logvar
|
| '''
|
|
|
| def __init__(self, latent_dim=256):
|
| super(VAE_Encoder, self).__init__()
|
| self.resnet = models.resnet18(pretrained=True)
|
| self.resnet.avgpool = nn.AdaptiveAvgPool2d((1, 1))
|
| self.resnet = nn.Sequential(
|
| *list(self.resnet.children())[:-1],
|
| FlattenLayer()
|
| )
|
|
|
| self.l_mu = nn.Linear(512, latent_dim)
|
| self.l_var = nn.Linear(512, latent_dim)
|
|
|
| def encode(self, x):
|
| hidden = self.resnet(x)
|
| mu = self.l_mu(hidden)
|
| logvar = self.l_var(hidden)
|
| return mu, logvar
|
|
|
| def reparameterize(self, mu, logvar):
|
| if self.training:
|
| std = torch.exp(0.5*logvar)
|
| eps = torch.randn_like(std)
|
| return mu + eps*std
|
|
|
| else:
|
| return mu
|
|
|
| def forward(self, x):
|
| mu, logvar = self.encode(x)
|
| z = self.reparameterize(mu, logvar)
|
| return z, mu, logvar
|
|
|
|
|
| class VAE_Decoder(nn.Module):
|
| '''
|
| VAE_Decoder: Decode noise to image
|
| '''
|
|
|
| def __init__(self, latent_dim, output_dim=3):
|
| super(VAE_Decoder, self).__init__()
|
| self.convs = nn.Sequential(
|
| UnflattenLayer(width=1),
|
| nn.ConvTranspose2d(latent_dim, 512, 4, 1, 0, bias=False),
|
| nn.ReLU(inplace=True),
|
| nn.ConvTranspose2d(512, 384, 4, 2, 1, bias=False),
|
| nn.BatchNorm2d(384),
|
| nn.ReLU(inplace=True),
|
| nn.ConvTranspose2d(384, 192, 4, 2, 1, bias=False),
|
| nn.BatchNorm2d(192),
|
| nn.ReLU(inplace=True),
|
| nn.ConvTranspose2d(192, 96, 4, 2, 1, bias=False),
|
| nn.BatchNorm2d(96),
|
| nn.ReLU(inplace=True),
|
| nn.ConvTranspose2d(96, 64, 4, 2, 1, bias=False),
|
| nn.BatchNorm2d(64),
|
| nn.ReLU(inplace=True),
|
| nn.ConvTranspose2d(64, 32, 4, 2, 1, bias=False),
|
| nn.BatchNorm2d(32),
|
| nn.LeakyReLU(inplace=True),
|
| nn.ConvTranspose2d(32, 3, 4, 2, 1, bias=False),
|
| nn.Tanh()
|
| )
|
|
|
| def forward(self, z):
|
| return self.convs(z)
|
|
|
| class ImageAE(nn.Module):
|
|
|
| def __init__(self):
|
| super(ImageAE, self).__init__()
|
| latent_dim = 512
|
| self.enc = VAE_Encoder(latent_dim)
|
| self.dec = VAE_Decoder(latent_dim)
|
|
|
| def forward(self, x):
|
| z, *_ = self.enc(x)
|
| out = self.dec(z)
|
|
|
| return out
|
|
|
| def load_ckpt(self, enc_path, dec_path):
|
| self.enc.load_state_dict(torch.load(enc_path, map_location='cpu'))
|
| self.dec.load_state_dict(torch.load(dec_path, map_location='cpu'))
|
|
|
|
|
| def get_pretraiend_ae(enc_path='pretrained/ae/vae/enc.pth', dec_path='pretrained/ae/vae/dec1.pth'):
|
| ae = ImageAE()
|
| ae.load_ckpt(enc_path, dec_path)
|
| print('load image auto-encoder')
|
| ae.eval()
|
| return ae
|
|
|
|
|
| def get_pretraiend_unet(path='pretrained/ae/unet/ckpt_srm.pth'):
|
| unet = UnetGenerator(3, 3, 8)
|
| unet.load_state_dict(torch.load(path, map_location='cpu'))
|
| print('load Unet')
|
| unet.eval()
|
| return unet
|
|
|
| if __name__ == "__main__":
|
| ae = get_pretraiend_ae()
|
| print(ae)
|
|
|
