Spaces:
Runtime error
Runtime error
| import torch.nn as nn | |
| import torch | |
| class Generator(torch.nn.Module): | |
| def __init__(self, nc_input=1, nc_output=1, ndf=128, nz=128, ngf=128, dropout_rate = 0.5 ): | |
| super(Generator, self).__init__() | |
| self.encoder = nn.Sequential( | |
| nn.Dropout(0.05), | |
| # input is (nc) x 64 x 64 | |
| nn.Conv2d(nc_input, ndf, 4, 2, 1, bias=False), | |
| nn.LeakyReLU(0.2, inplace=True), | |
| # state size. (ndf) x 32 x 32 | |
| nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False), | |
| nn.BatchNorm2d(ndf * 2), | |
| nn.LeakyReLU(0.2, inplace=True), | |
| # state size. (ndf*2) x 16 x 16 | |
| nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False), | |
| nn.BatchNorm2d(ndf * 4), | |
| nn.LeakyReLU(0.2, inplace=True), | |
| # state size. (ndf*4) x 8 x 8 | |
| nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False), | |
| nn.BatchNorm2d(ndf * 8), | |
| nn.LeakyReLU(0.2, inplace=True), | |
| # state size. (ndf*8) x 4 x 4 | |
| nn.Conv2d(ndf * 8, 1, 1, 1, 0, bias=False), | |
| ##nn.Conv2d(1, 1, 5, 1, 0, bias=False), | |
| ##nn.Sigmoid() | |
| ) | |
| self.linearEncoder = nn.Sequential( | |
| nn.Linear(64, 128) | |
| ) | |
| self.decoder = nn.Sequential( | |
| # input is Z, going into a convolution | |
| nn.Dropout(0.05), | |
| nn.ConvTranspose2d(nz, ngf * 8, 4, 1, 0, bias=False), | |
| nn.BatchNorm2d(ngf * 8), | |
| nn.ReLU(True), | |
| nn.Dropout(dropout_rate), | |
| # state size. (ngf*8) x 4 x 4 == 1024 x 4 x 4 | |
| nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False), | |
| nn.BatchNorm2d(ngf * 4), | |
| nn.ReLU(True), | |
| nn.Dropout(dropout_rate), # state size. (ngf*4) x 8 x 8 == 512 x 4 x 4 | |
| nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False), | |
| nn.BatchNorm2d(ngf * 2), | |
| nn.ReLU(True), | |
| nn.Dropout(dropout_rate), # state size. (ngf*2) x 16 x 16 == 256 x 4 x 4 | |
| nn.ConvTranspose2d(ngf * 2, ngf, 4, 2, 1, bias=False), | |
| nn.BatchNorm2d(ngf), | |
| nn.ReLU(True), | |
| nn.Dropout(dropout_rate), # state size. (ngf) x 32 x 32 == 128 x 4 x 4 | |
| nn.ConvTranspose2d(ngf, ngf, 4, 2, 1, bias=False), | |
| nn.BatchNorm2d(ngf), | |
| nn.ReLU(True), | |
| nn.ConvTranspose2d( ngf, nc_output, 4, 2, 1, bias=False), | |
| nn.Sigmoid() | |
| # state size. (nc) x 64 x 64 | |
| ) | |
| def forward(self, x): | |
| encoded = self.forward_encoder(x) | |
| decoded = self.forward_decoder(encoded) | |
| return decoded | |
| def forward_encoder(self, x): | |
| encoded = self.encoder(x).reshape(-1, 64) | |
| return self.linearEncoder(encoded).unsqueeze(2).unsqueeze(2) | |
| def forward_decoder(self, encoded): | |
| decoded = self.decoder(encoded) | |
| return decoded | |