pytorch_WGAN-gp

import torch
import torchvision
from torchvision import transforms
from torch.utils.data import DataLoader
from torch import nn
import torch.autograd as autograd
import matplotlib.pyplot as plt
import os
import numpy as np
import matplotlib

def to_img(x):
    x = 0.5 * (x + 1)
    x = x.clamp(0, 1)
    x = x.view(x.size(0),1,28,28)
    return x

def imshow(img,filename=None):
    npimg = img.numpy()
    plt.axis('off')
    array = np.transpose(npimg, (1, 2, 0))    
    if filename!=None:
        matplotlib.image.imsave(filename, array)       
    else:
        plt.imshow(array  ) 
#        plt.savefig(filename) 保存图片
        plt.show()

data_dir = './fashion_mnist/'
train_dataset = torchvision.datasets.FashionMNIST(data_dir, train=True, 
                                                  transform=img_transform,download=True)
#train_loader = DataLoader(train_dataset,batch_size=1024, shuffle=True)
train_loader = DataLoader(train_dataset,batch_size=256, shuffle=True)

val_dataset = torchvision.datasets.FashionMNIST(data_dir, train=False, 
                                                transform=img_transform)
test_loader = DataLoader(val_dataset, batch_size=10, shuffle=False)
#指定设备
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)

class WGAN_D(nn.Module):
    def __init__(self,inputch=1):
        super(WGAN_D, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(inputch, 64,4, 2, 1),  # batch, 64, 28, 28
            nn.LeakyReLU(0.2, True),
            nn.InstanceNorm2d(64, affine=True)   )
        self.conv2 = nn.Sequential(
            nn.Conv2d(64, 128,4, 2, 1),  # batch, 64, 14, 14
            nn.LeakyReLU(0.2, True),
            nn.InstanceNorm2d(128, affine=True)   )
        self.fc = nn.Sequential(
            nn.Linear(128*7*7, 1024),
            nn.LeakyReLU(0.2, True),            )
        self.fc2 =nn.Sequential(
                nn.InstanceNorm1d(1, affine=True),
                nn.Flatten(),
                nn.Linear(1024, 1)  )
        

               
    def forward(self, x,*arg):#batch, width, height, channel=1
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        x = x.reshape(x.size(0),1, -1)        
        x = self.fc2(x)
        return x.view(-1, 1).squeeze(1)

class WGAN_G(nn.Module):
    def __init__(self, input_size,input_n=1):
        super(WGAN_G, self).__init__()

        self.fc1 = nn.Sequential(
            nn.Linear(input_size*input_n, 1024),
            nn.ReLU(True),
            nn.BatchNorm1d(1024)  )

        self.fc2 = nn.Sequential(
            nn.Linear(1024,7*7*128),
            nn.ReLU(True),
            nn.BatchNorm1d(7*7*128)   )
        self.upsample1 = nn.Sequential(
            nn.ConvTranspose2d(128, 64, 4, 2, padding=1, bias=False),  # batch, 64, 14, 14
            nn.ReLU(True),
            nn.BatchNorm2d(64)   )
        self.upsample2 = nn.Sequential(
            nn.ConvTranspose2d(64, 1, 4, 2, padding=1, bias=False),  # batch, 64, 28, 28
            nn.Tanh(),  )
        

    def forward(self, x,*arg):
        x = self.fc1(x)
        x = self.fc2(x)
        
        x = x.view(x.size(0), 128, 7, 7)
        x = self.upsample1(x)
        img = self.upsample2(x)

        return img

# Loss weight for gradient penalty
lambda_gp = 10
def compute_gradient_penalty(D, real_samples, fake_samples,y_one_hot):

    eps = torch.FloatTensor(real_samples.size(0),1,1,1).uniform_(0,1).to(device)
    # Get random interpolation between real and fake samples
    X_inter = (eps * real_samples + ((1 - eps) * fake_samples)).requires_grad_(True)
    d_interpolates = D(X_inter,y_one_hot)
    fake = torch.full((real_samples.size(0), ), 1, device=device)
    
    # Get gradient
    gradients = autograd.grad( outputs=d_interpolates,
            inputs=X_inter,
            grad_outputs=fake,
            create_graph=True,
            retain_graph=True,
            only_inputs=True,
    )[0]
    gradients = gradients.view(gradients.size(0), -1)
    gradient_penaltys = ((gradients.norm(2, dim=1) - 1) ** 2).mean() * lambda_gp
    return gradient_penaltys

def train(D,G,outdir,z_dimension ,num_epochs = 30):
    d_optimizer = torch.optim.Adam(D.parameters(), lr=0.001)
    g_optimizer = torch.optim.Adam(G.parameters(), lr=0.001)
    
    os.makedirs(outdir, exist_ok=True)

    # train
    for epoch in range(num_epochs):
        for i, (img, lab) in enumerate(train_loader):
            num_img = img.size(0)
            # =================train discriminator

            real_img = img.to(device)
            y_one_hot = torch.zeros(lab.shape[0],10).scatter_(1,
                                   lab.view(lab.shape[0],1),1).to(device)
            for ii in range(5):
                d_optimizer.zero_grad()
        
                # compute loss of real_img
                real_out = D(real_img,y_one_hot)# closer to 1 means better
                # compute loss of fake_img
                z = torch.randn(num_img, z_dimension).to(device)
                fake_img = G(z,y_one_hot)
                fake_out = D(fake_img,y_one_hot)# closer to 0 means better
                
                gradient_penalty = compute_gradient_penalty(D, 
                                        real_img.data, fake_img.data,y_one_hot)
    
                # Loss measures generator's ability to fool the discriminator
                d_loss = -torch.mean(real_out) + torch.mean(fake_out) + gradient_penalty
                d_loss.backward()
                d_optimizer.step()
    
            # ===============train generator
            # compute loss of fake_img
            for ii in range(1):
                g_optimizer.zero_grad()
                z = torch.randn(num_img, z_dimension).to(device)
                fake_img = G(z,y_one_hot)
                fake_out = D(fake_img,y_one_hot)
                g_loss = -torch.mean(fake_out)
                g_loss.backward()
                g_optimizer.step()
                
        fake_images = to_img(fake_img.cpu().data)
        real_images = to_img(real_img.cpu().data)
        rel = torch.cat([to_img(real_images[:10]),fake_images[:10]],axis = 0)
        imshow(torchvision.utils.make_grid(rel,nrow=10),
              os.path.join(outdir, 'fake_images-{}.png'.format(epoch+1) ) )

        
        print('Epoch [{}/{}], d_loss: {:.6f}, g_loss: {:.6f} '
                      'D real: {:.6f}, D fake: {:.6f}'
                      .format(epoch, num_epochs, d_loss.data, g_loss.data,
                              real_out.data.mean(), fake_out.data.mean()))
        
    torch.save(G.state_dict(), os.path.join(outdir, 'generator.pth'  ) )
    torch.save(D.state_dict(), os.path.join(outdir, 'discriminator.pth'  ) )   

def displayAndTest(D,G,z_dimension):
    # 可视化结果
    sample = iter(test_loader)
    #|images, labels = sample.next()
    images,labels = next(sample)
    y_one_hot = torch.zeros(labels.shape[0],10).scatter_(1,
                                   labels.view(labels.shape[0],1),1).to(device)
    
    num_img = images.size(0)
    with torch.no_grad():
        z = torch.randn(num_img, z_dimension).to(device)
        fake_img = G(z,y_one_hot)
    fake_images = to_img(fake_img.cpu().data)
    rel = torch.cat([to_img(images[:10]),fake_images[:10]],axis = 0)
    imshow(torchvision.utils.make_grid(rel,nrow=10))
    print(labels[:10])     

if __name__ == '__main__': 
    
    z_dimension = 40  # noise dimension
    
    D = WGAN_D().to(device)  # discriminator model
    G = WGAN_G(z_dimension).to(device)  # generator model 
    train(D,G,'./w_img',z_dimension)
    displayAndTest(D,G,z_dimension)