WGANGP not working as expected

[ncuxomun]

Hi, can you please tell me if I am doing something wrong here, especially the manual update for the generator and the critic:

class Generator(nn.Module):
    def __init__(self, latent_dim=64, img_shape=None):
        super().__init__()
        self.img_shape = img_shape
        self.init_size = 8 #self.img_shape[1] // 4

        self.l1 = nn.Sequential(
            nn.Linear(latent_dim, 64*self.init_size**2), nn.LeakyReLU(0.2, inplace=True))
        self.conv_blocks = nn.Sequential(
            nn.BatchNorm2d(64),
            nn.Upsample(scale_factor=2),
            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, padding=0),
            nn.BatchNorm2d(64),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Upsample(scale_factor=2),
            nn.Conv2d(in_channels=64, out_channels=32, kernel_size=3, padding=0),
            nn.BatchNorm2d(32),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Upsample(scale_factor=2),
            nn.Conv2d(in_channels=32, out_channels=16, kernel_size=3, padding=0),
            nn.BatchNorm2d(16),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Upsample(scale_factor=2),
            nn.Conv2d(in_channels=16, out_channels=8, kernel_size=3, padding=1),
            nn.BatchNorm2d(8),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(in_channels=8, out_channels=img_shape[0], kernel_size=3, padding=1),
            nn.Tanh()
        )
    
    def forward(self, z):
        out = self.l1(z)
        out = out.view(out.shape[0], 64, self.init_size, self.init_size)
        img = self.conv_blocks(out)
        return img

class Critic(nn.Module):
    def __init__(self, img_shape):
        super().__init__()
        self.disc = nn.Sequential(
            nn.Conv2d(in_channels=img_shape[0], out_channels=16, kernel_size=4, stride=2),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(16, 32, kernel_size=4, stride=2),
            nn.BatchNorm2d(32),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.BatchNorm2d(64),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(64, 128, kernel_size=4, stride=2),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2, inplace=True),
        )

        # The height and width of downsampled image
        #
        ds_size = 2 ** 4

        self.adv_layer = nn.Sequential(nn.Linear(128 * ds_size, 1))
    def forward(self, img):
        out = self.disc(img)
        # import pdb; pdb.set_trace()
        out = out.view(out.shape[0], -1)
        validity = self.adv_layer(out)
        return validity       

class WGANGP(pl.LightningModule):
    def __init__(self, latent_dim=128, lr=0.0002, lambda_pen=10, crit_repeats=5):
        super().__init__()
        self.save_hyperparameters()
        self.latent_dim = latent_dim
        self.lr = lr
        self.lambda_pen = lambda_pen
        self.crit_repeats = crit_repeats
        self.b1 = 0.0
        self.b2 = 0.9

        ### initializing networks
        img_shape = (1, 100, 100)
        self.generator = Generator(self.latent_dim, img_shape)
        self.critic = Critic(img_shape)

        # application of weight
        self.generator.apply(self.weights_init)
        self.critic.apply(self.weights_init)
        #
        self.validation_z = torch.randn(10, self.latent_dim)
        self.example_input_array = torch.zeros(10, self.latent_dim)

        # Important: This property activates manual optimization.
        self.automatic_optimization = False # True - Auto // # False - Manual update

    def forward(self, z):
        return self.generator(z)

    ### weight initialization
    def weights_init(self, m):
        if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
            torch.nn.init.normal_(m.weight, 0.0, 0.02)
        if isinstance(m, nn.BatchNorm2d):
            torch.nn.init.normal_(m.weight, 0.0, 0.02)
            torch.nn.init.constant_(m.bias, 0)
        if isinstance(m, nn.Linear):
            torch.nn.init.normal_(m.weight, 0.0, 0.02)
            torch.nn.init.constant_(m.bias, 0)
        
    def training_step(self, batch, batch_idx):
        imgs = batch

        # # sample noise
        # z = torch.randn(imgs.shape[0], self.latent_dim)
        # z = z.type_as(imgs)

        # optimizers, manual access
        g_opt, c_opt = self.optimizers()

        # update critic
        mean_iteration_critic_loss = 0
        for _ in range(self.crit_repeats):
            c_opt.zero_grad()
            # sample noise
            z = torch.randn(imgs.shape[0], self.latent_dim).type_as(imgs)
            # fake image
            fake = self(z)
            crit_fake_pred = self.critic(fake.detach())
            crit_real_pred = self.critic(imgs)
            # eps
            epsilon = torch.rand(len(imgs), 1, 1, 1, device=self.device, requires_grad=True)
            # gradient penalty
            gp = self.gradient_penalty(self.critic, imgs, fake, epsilon)

            # critic loss
            critic_loss = torch.mean(crit_fake_pred) - torch.mean(crit_real_pred) + self.lambda_pen * gp

            # Keep track of the average critic loss in this batch
            mean_iteration_critic_loss += critic_loss.item() / crit_repeats

            # Update gradients
            self.manual_backward(critic_loss)
            # Update optimizer
            c_opt.step()

        # log critic average loss
        self.log('c_loss_mean', mean_iteration_critic_loss, prog_bar=True)
        
        # update generator
        g_opt.zero_grad()
        # sample new noise
        z_new = torch.randn(imgs.shape[0], self.latent_dim).type_as(imgs)
        # new fake image
        fake_new = self(z_new)
        crit_fake_pred = self.critic(fake_new)
        # generator loss
        gen_loss = -torch.mean(crit_fake_pred)

        # Update gradients
        self.manual_backward(gen_loss)
        # Update optimizer
        g_opt.step()

        # log generator average loss
        self.log('g_loss', gen_loss, prog_bar=True)

    def gradient_penalty(self, crit, real, fake, epsilon):
        # mix/interpolate images
        mixed_images = real * epsilon + fake * (1 - epsilon)

        # Calculate the critic's scores on the mixed images
        mixed_scores = crit(mixed_images)

        # Take the gradient of the scores with respect to the images
        gradient = torch.autograd.grad(
            inputs=mixed_images,
            outputs=mixed_scores,
            grad_outputs=torch.ones_like(mixed_scores),
            create_graph=True,
            retain_graph=True,
        )[0]

        # Flatten the gradients so that each row captures one image
        gradient = gradient.view(len(gradient), -1)

        # Calculate the magnitude of every row
        gradient_norm = gradient.norm(2, dim=1)

        # Penalize the mean squared distance of the gradient norms from 1
        gradient_penalty = torch.mean((gradient_norm - 1) ** 2)

        return gradient_penalty
    
    def configure_optimizers(self):
        opt_g = torch.optim.Adam(self.generator.parameters(), lr=self.lr, betas=(self.b1, self.b2))
        opt_c = torch.optim.Adam(self.critic.parameters(), lr=self.lr, betas=(self.b1, self.b2))
        return opt_g, opt_c

    def on_epoch_end(self):
        z = self.validation_z.to(self.device)
        # log sampled images
        sample_imgs = self(z)
        grid = torchvision.utils.make_grid(sample_imgs)
        self.logger.experiment.add_image('generated_images', grid, self.current_epoch)

# defining the hyperparameters
n_epochs = 1000
z_dim = 50
batch_size = 64
lr = 0.0002
c_lambda = 10
crit_repeats = 5

desired output
current output

SOLVED

[ncuxomun]

Solved.

[frankNiessen]

how?

[LebinDing749]

How?