| import torch |
| import torch.nn as nn |
|
|
| class ConvGenerator(nn.Module): |
| def __init__(self, latent_dim=100, channels=1): |
| super(ConvGenerator, self).__init__() |
| self.latent_dim = latent_dim |
| |
| self.init_size = 7 |
| self.l1 = nn.Sequential(nn.Linear(latent_dim, 128 * self.init_size ** 2)) |
|
|
| self.conv_blocks = nn.Sequential( |
| nn.BatchNorm2d(128), |
| nn.Upsample(scale_factor=2), |
| nn.Conv2d(128, 128, 3, stride=1, padding=1), |
| nn.BatchNorm2d(128, 0.8), |
| nn.LeakyReLU(0.2, inplace=True), |
| nn.Upsample(scale_factor=2), |
| nn.Conv2d(128, 64, 3, stride=1, padding=1), |
| nn.BatchNorm2d(64, 0.8), |
| nn.LeakyReLU(0.2, inplace=True), |
| nn.Conv2d(64, channels, 3, stride=1, padding=1), |
| nn.Tanh() |
| ) |
|
|
| def forward(self, z): |
| out = self.l1(z) |
| out = out.view(out.shape[0], 128, self.init_size, self.init_size) |
| img = self.conv_blocks(out) |
| return img |
|
|
| class ConvDiscriminator(nn.Module): |
| def __init__(self, channels=1): |
| super(ConvDiscriminator, self).__init__() |
|
|
| def discriminator_block(in_filters, out_filters, bn=True): |
| block = [nn.Conv2d(in_filters, out_filters, 3, 2, 1), |
| nn.LeakyReLU(0.2, inplace=True), |
| nn.Dropout2d(0.25)] |
| if bn: |
| block.append(nn.BatchNorm2d(out_filters, 0.8)) |
| return block |
|
|
| self.model = nn.Sequential( |
| *discriminator_block(channels, 16, bn=False), |
| *discriminator_block(16, 32), |
| *discriminator_block(32, 64), |
| *discriminator_block(64, 128), |
| ) |
|
|
| |
| ds_size = 28 // 2**4 |
| self.adv_layer = nn.Sequential( |
| nn.Linear(128 * ds_size ** 2, 1), |
| nn.Sigmoid() |
| ) |
|
|
| def forward(self, img): |
| out = self.model(img) |
| out = out.view(out.shape[0], -1) |
| validity = self.adv_layer(out) |
| return validity |
