|
import torch |
|
import torch.nn as nn |
|
|
|
class Generator(nn.Module): |
|
def __init__(self, z_dim=100, img_channels=3): |
|
super(Generator, self).__init__() |
|
self.gen = nn.Sequential( |
|
|
|
nn.ConvTranspose2d(z_dim, 512, 4, 1, 0, bias=False), |
|
nn.BatchNorm2d(512), |
|
nn.ReLU(True), |
|
|
|
nn.ConvTranspose2d(512, 256, 4, 2, 1, bias=False), |
|
nn.BatchNorm2d(256), |
|
nn.ReLU(True), |
|
|
|
nn.ConvTranspose2d(256, 128, 4, 2, 1, bias=False), |
|
nn.BatchNorm2d(128), |
|
nn.ReLU(True), |
|
|
|
nn.ConvTranspose2d(128, 64, 4, 2, 1, bias=False), |
|
nn.BatchNorm2d(64), |
|
nn.ReLU(True), |
|
|
|
nn.ConvTranspose2d(64, img_channels, 4, 2, 1, bias=False), |
|
nn.Tanh() |
|
|
|
) |
|
|
|
def forward(self, input): |
|
return self.gen(input) |
|
|
|
class Discriminator(nn.Module): |
|
def __init__(self, img_channels=3): |
|
super(Discriminator, self).__init__() |
|
self.disc = nn.Sequential( |
|
|
|
nn.Conv2d(img_channels, 64, 4, 2, 1, bias=False), |
|
nn.LeakyReLU(0.2, inplace=True), |
|
|
|
nn.Conv2d(64, 128, 4, 2, 1, bias=False), |
|
nn.BatchNorm2d(128), |
|
nn.LeakyReLU(0.2, inplace=True), |
|
|
|
nn.Conv2d(128, 256, 4, 2, 1, bias=False), |
|
nn.BatchNorm2d(256), |
|
nn.LeakyReLU(0.2, inplace=True), |
|
|
|
nn.Conv2d(256, 512, 4, 2, 1, bias=False), |
|
nn.BatchNorm2d(512), |
|
nn.LeakyReLU(0.2, inplace=True), |
|
|
|
nn.Conv2d(512, 1, 4, 1, 0, bias=False), |
|
nn.Sigmoid() |
|
) |
|
|
|
def forward(self, input): |
|
return self.disc(input).view(-1, 1).squeeze(1) |
|
|
|
batch_size = 32 |
|
latent_vector_size = 100 |
|
|
|
generator = Generator() |
|
discriminator = Discriminator() |
|
|
|
generator.load_state_dict(torch.load('netG.pth', map_location=torch.device('cpu') )) |
|
discriminator.load_state_dict(torch.load('netD.pth', map_location=torch.device('cpu') )) |
|
|