Here are two GAN in pytorch that are pretty simple and easy to follow if they help you.

  
      
      GitHub
  
  
    

MatthewR2D2/Pytorch 7

This repository is for my learning Pytorch. Contribute to MatthewR2D2/Pytorch development by creating an account on GitHub.


  
  
    
    
  
  


Here is how to set up two models GEN and DESCRIM and train them
# Set models for training
Disc.train()
Gen.train()
for epoch in range(num_epochs):
    # Each batch
    for batch_i, (real_images, _) in enumerate(train_loader):
        batch_size = real_images.size(0)
        
        ## Important rescaling step ## 
        # rescale input images from [0,1) to [-1, 1)
        real_images = real_images*2 - 1
        # Train Discriminator
        d_optimizer.zero_grad()
        # Train with real images
        D_real = Disc(real_images)
        d_real_loss = real_loss(D_real, smooth=True) # Use label smoothing
        
        # Next step train with fake images
        # Generate fake images
        z = np.random.uniform(-1, 1, size=(batch_size, z_size)) # Random Noise
        z = torch.from_numpy(z).float() # Convert to a flot
        fake_images = Gen(z) # Predict with Generator (DO NOT TRAIN GENERATOR HERE) train one at a time
        
        # Computer fake loss
        D_fake = Disc(fake_images)
        d_fake_loss = fake_loss(D_fake)
        # add up loss and back prop
        d_loss = d_real_loss + d_fake_loss
        d_loss.backward()
        d_optimizer.step()
        
        '''
        # Generator training
        '''
        g_optimizer.zero_grad()
        # Generator fake images and train
        z = np.random.uniform(-1, 1, size=(batch_size, z_size))
        z = torch.from_numpy(z).float()
        fake_images = Gen(z)
        
        # Compute the discriminator losses on fake images 
        # using flipped labels!
        D_fake = Disc(fake_images)
        g_loss = real_loss(D_fake) # use real loss to flip labels
        
        # perform backprop
        g_loss.backward()
        g_optimizer.step()
        
        # Print some loss stats
        if batch_i % print_every == 0:
            # print discriminator and generator loss
            print('Epoch [{:5d}/{:5d}] | d_loss: {:6.4f} | g_loss: {:6.4f}'.format(
                    epoch+1, num_epochs, d_loss.item(), g_loss.item()))
    
    ## AFTER EACH EPOCH##
    # append discriminator loss and generator loss
    losses.append((d_loss.item(), g_loss.item()))
    
    # generate and save sample, fake images
    Gen.eval() # eval mode for generating samples
    samples_z = Gen(fixed_z)
    samples.append(samples_z)
    Gen.train() # back to train mode