Spaces:
Runtime error
Runtime error
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import torch.optim as optim | |
| from torch.nn.parallel.distributed import DistributedDataParallel as DDP | |
| from .nn import mean_flat | |
| #from . import dist_util | |
| import functools | |
| class AdversarialLoss(nn.Module): | |
| def __init__(self, gan_type='WGAN_GP', gan_k=1, | |
| lr_dis=1e-5 ): | |
| super(AdversarialLoss, self).__init__() | |
| self.gan_type = gan_type | |
| self.gan_k = gan_k | |
| model = NLayerDiscriminator().cuda() | |
| self.discriminator = DDP( | |
| model, | |
| device_ids=[torch.device('cuda')], | |
| output_device=torch.device('cuda'), | |
| broadcast_buffers=False, | |
| bucket_cap_mb=128, | |
| find_unused_parameters=False, | |
| ) | |
| if (gan_type in ['WGAN_GP', 'GAN']): | |
| self.optimizer = optim.Adam( | |
| self.discriminator.parameters(), | |
| lr=lr_dis | |
| ) | |
| def forward(self, fake, real): | |
| fake_detach = fake.detach() | |
| for _ in range(self.gan_k): | |
| self.optimizer.zero_grad() | |
| d_fake = self.discriminator(fake_detach) | |
| d_real = self.discriminator(real) | |
| if (self.gan_type.find('WGAN') >= 0): | |
| loss_d = (d_fake - d_real).mean() | |
| if self.gan_type.find('GP') >= 0: | |
| epsilon = torch.rand(real.size(0), 1, 1, 1).cuda() | |
| epsilon = epsilon.expand(real.size()) | |
| hat = fake_detach.mul(1 - epsilon) + real.mul(epsilon) | |
| hat.requires_grad = True | |
| d_hat = self.discriminator(hat) | |
| gradients = torch.autograd.grad( | |
| outputs=d_hat.sum(), inputs=hat, | |
| retain_graph=True, create_graph=True, only_inputs=True | |
| )[0] | |
| gradients = gradients.view(gradients.size(0), -1) | |
| gradient_norm = gradients.norm(2, dim=1) | |
| gradient_penalty = 10 * gradient_norm.sub(1).pow(2).mean() | |
| loss_d += gradient_penalty | |
| # print('d loss:', loss_d) | |
| # Discriminator update | |
| loss_d.backward() | |
| self.optimizer.step() | |
| d_fake_for_g = self.discriminator(fake) | |
| if (self.gan_type.find('WGAN') >= 0): | |
| loss_g = -d_fake_for_g | |
| # Generator loss | |
| return mean_flat(loss_g) | |
| def conv3x3(in_channels, out_channels, stride=1): | |
| return nn.Conv2d(in_channels, out_channels, kernel_size=3, | |
| stride=stride, padding=1, bias=True) | |
| def conv7x7(in_channels, out_channels, stride=1): | |
| return nn.Conv2d(in_channels, out_channels, kernel_size=7, | |
| stride=stride, padding=3, bias=True) | |
| class Discriminator(nn.Module): | |
| def __init__(self, ): | |
| super(Discriminator, self).__init__() | |
| self.conv1 = conv7x7(3, 32) | |
| self.norm1 = nn.InstanceNorm2d(32, affine=True) | |
| self.LReLU1 = nn.LeakyReLU(0.2) | |
| self.conv2 = conv3x3(32, 32, 2) | |
| self.norm2 = nn.InstanceNorm2d(32, affine=True) | |
| self.LReLU2 = nn.LeakyReLU(0.2) | |
| self.conv3 = conv3x3(32, 64) | |
| self.norm3 = nn.InstanceNorm2d(64, affine=True) | |
| self.LReLU3 = nn.LeakyReLU(0.2) | |
| self.conv4 = conv3x3(64, 64, 2) | |
| self.norm4 = nn.InstanceNorm2d(64, affine=True) | |
| self.LReLU4 = nn.LeakyReLU(0.2) | |
| self.conv5 = conv3x3(64, 128) | |
| self.norm5 = nn.InstanceNorm2d(128, affine=True) | |
| self.LReLU5 = nn.LeakyReLU(0.2) | |
| self.conv6 = conv3x3(128, 128, 2) | |
| self.norm6 = nn.InstanceNorm2d(128, affine=True) | |
| self.LReLU6 = nn.LeakyReLU(0.2) | |
| self.conv7 = conv3x3(128, 256) | |
| self.norm7 = nn.InstanceNorm2d(256, affine=True) | |
| self.LReLU7 = nn.LeakyReLU(0.2) | |
| self.conv8 = conv3x3(256, 256, 2) | |
| self.norm8 = nn.InstanceNorm2d(256, affine=True) | |
| self.LReLU8 = nn.LeakyReLU(0.2) | |
| self.conv9 = conv3x3(256, 512) | |
| self.norm9 = nn.InstanceNorm2d(512, affine=True) | |
| self.LReLU9 = nn.LeakyReLU(0.2) | |
| self.conv10 = conv3x3(512, 512, 2) | |
| self.norm10 = nn.InstanceNorm2d(512, affine=True) | |
| self.LReLU10 = nn.LeakyReLU(0.2) | |
| self.conv11 = conv3x3(512, 32) | |
| self.norm11 = nn.InstanceNorm2d(32, affine=True) | |
| self.LReLU11 = nn.LeakyReLU(0.2) | |
| self.conv12 = conv3x3(32, 1) | |
| def forward(self, x): | |
| x = self.LReLU1(self.norm1(self.conv1(x))) | |
| x = self.LReLU2(self.norm2(self.conv2(x))) | |
| x = self.LReLU3(self.norm3(self.conv3(x))) | |
| x = self.LReLU4(self.norm4(self.conv4(x))) | |
| x = self.LReLU5(self.norm5(self.conv5(x))) | |
| x = self.LReLU6(self.norm6(self.conv6(x))) | |
| x = self.LReLU7(self.norm7(self.conv7(x))) | |
| x = self.LReLU8(self.norm8(self.conv8(x))) | |
| x = self.LReLU9(self.norm9(self.conv9(x))) | |
| x = self.LReLU10(self.norm10(self.conv10(x))) | |
| x = self.LReLU11(self.norm11(self.conv11(x))) | |
| x = self.conv12(x) | |
| return x | |
| def get_norm_layer(norm_type='instance'): | |
| """Return a normalization layer | |
| Parameters: | |
| norm_type (str) -- the name of the normalization layer: batch | instance | none | |
| For BatchNorm, we use learnable affine parameters and track running statistics (mean/stddev). | |
| For InstanceNorm, we do not use learnable affine parameters. We do not track running statistics. | |
| """ | |
| if norm_type == 'batch': | |
| norm_layer = functools.partial(nn.BatchNorm2d, affine=True, track_running_stats=True) | |
| elif norm_type == 'instance': | |
| norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False) | |
| elif norm_type == 'none': | |
| def norm_layer(x): return Identity() | |
| else: | |
| raise NotImplementedError('normalization layer [%s] is not found' % norm_type) | |
| return norm_layer | |
| class NLayerDiscriminator(nn.Module): | |
| """Defines a PatchGAN discriminator""" | |
| def __init__(self, input_nc=3, ndf=64, n_layers=3 ): | |
| """Construct a PatchGAN discriminator | |
| Parameters: | |
| input_nc (int) -- the number of channels in input images | |
| ndf (int) -- the number of filters in the last conv layer | |
| n_layers (int) -- the number of conv layers in the discriminator | |
| norm_layer -- normalization layer | |
| """ | |
| super(NLayerDiscriminator, self).__init__() | |
| norm_layer = get_norm_layer(norm_type='instance') | |
| if type(norm_layer) == functools.partial: # no need to use bias as BatchNorm2d has affine parameters | |
| use_bias = norm_layer.func == nn.InstanceNorm2d | |
| else: | |
| use_bias = norm_layer == nn.InstanceNorm2d | |
| kw = 4 | |
| padw = 1 | |
| sequence = [nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)] | |
| nf_mult = 1 | |
| nf_mult_prev = 1 | |
| for n in range(1, n_layers): # gradually increase the number of filters | |
| nf_mult_prev = nf_mult | |
| nf_mult = min(2 ** n, 8) | |
| sequence += [ | |
| nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw, bias=use_bias), | |
| norm_layer(ndf * nf_mult), | |
| nn.LeakyReLU(0.2, True) | |
| ] | |
| nf_mult_prev = nf_mult | |
| nf_mult = min(2 ** n_layers, 8) | |
| sequence += [ | |
| nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias), | |
| norm_layer(ndf * nf_mult), | |
| nn.LeakyReLU(0.2, True) | |
| ] | |
| sequence += [nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)] # output 1 channel prediction map | |
| self.model = nn.Sequential(*sequence) | |
| def forward(self, input): | |
| """Standard forward.""" | |
| return self.model(input) |