GAN实战

import torch
from   torch import nn,optim,autograd
import numpy as np
import visdom
from   torch.nn import functional as F
from   matplotlib import pyplot as plt
import random

h_dim =400
batchsz=512
viz =visdom.Visdom()

class Generator(nn.Module):
    def __init__(self):
        super(Generator,self).__init__()

        self.net=nn.Sequential(
            nn.Linear(2,h_dim),
            nn.ReLU(),
            nn.Linear(h_dim,h_dim),
            nn.ReLU(),
            nn.Linear(h_dim,h_dim),
            nn.ReLU(),
            nn.Linear(h_dim,2)
        )
    def forward(self,z):

        output=self.net(z)
        return output
class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator,self).__init__()

        self.net=nn.Sequential(
            nn.Linear(2,h_dim),
            nn.ReLU(),
            nn.Linear(h_dim,h_dim),
            nn.ReLU(),
            nn.Linear(h_dim,h_dim),
            nn.ReLU(),
            nn.Linear(h_dim,1),
            nn.Sigmoid()
        )
    def forward(self,x):
        output=self.net(x)
        #[b,1]--->[b]
        return output.view(-1)

def data_generator():
    scale = 2
    centers =[
        (1,0),
        (-1,0),
        (0,1),
        (0,-1),
        (1./np.sqrt(2),1./np.sqrt(2)),
        (1./np.sqrt(2),-1./np.sqrt(2)),
        (-1./np.sqrt(2),1./np.sqrt(2)),
        (-1./np.sqrt(2),-1./np.sqrt(2)),

    ] #列表中添加tuple
    centers=[(scale*x,scale*y) for x,y in centers]
    while True:
        dataset = []
        for i in range(batchsz):
            point = np.random.randn(2) * .02
            center = random.choice(centers)
            point[0] += center[0]
            point[1] += center[1]
            dataset.append(point)
        dataset = np.array(dataset, dtype='float32')
        dataset /= 1.414  # stdev
        yield dataset

def generator_image(D,G,xr,epoch):
    N_POINTS =128
    RANGE =3

    plt.clf()##Clear figure清除所有轴，但是窗口打开，这样它可以被重复使用。
    points =np.zeros((N_POINTS, N_POINTS, 2), dtype='float32')
    points[:,:,0]=np.linspace(-RANGE,RANGE,N_POINTS)[:,None] #在行上进行在-range 到 range的线性插值shape :128*128       
    points[:,:,1]=np.linspace(-RANGE,RANGE,N_POINTS)[None,:] #在列上进行在-range 到 range的线性插值
    points =points.reshape((-1,2))  #points的shape(128*128*2)在里边的一维元素长度为2 ，reshape之后 (16384, 2)        哦
    # (16384, 2)
    # print('p:', points.shape)

    # draw contour
    with torch.no_grad():
        points=torch.Tensor(points).cuda()# [16384, 2]
        disc_map =D(points).cpu().numpy()# [16384]
    x=y=np.linspace(-RANGE,RANGE,N_POINTS)
    #添加等高线
    cs = plt.contour(x, y, disc_map.reshape((len(x), len(y))))
    plt.clabel(cs,inline=1,fontsize=10)#画等高线并标出函数值
    
    cs = plt.contour(x, y, disc_map.reshape((len(x), len(y))))
    # draw samples
    with torch.no_grad():
        z = torch.randn(batchsz, 2).cuda() # [b, 2]
        samples = G(z).cpu().numpy() # [b, 2]
    xr=xr.cpu().numpy() #GPU tensor不能直接转换成numpy先转换成CPU tensor，再转换成numpy
    plt.scatter(xr[:, 0], xr[:, 1], c='orange', marker='.')
    plt.scatter(samples[:, 0], samples[:, 1], c='green', marker='+')
    
    viz.matplot(plt,win='contour',opts=dict(title='p(x):%d'%epoch))
def gradient_penalty(D,xr,xf):
    
    LAMNDA=0.3
    #only constrait for Discriminator
    
    xf= xf.detach()
    xr=xr.detach()
    
    #[b,1]==>[b,2]
    alpha = torch.rand(batchsz,1).cuda()
    alpha=alpha.expand_as(xr)#对采样点一维数据有两个元素
    
    interpolates=alpha*xr+(1-alpha)*xf#在真假数据之间进行线性插值选择对应的样本
    interpolates.requires_grad_()
    
    disc_interpolates=D(interpolates)
    
    gradients=autograd.grad(outputs=disc_interpolates,inputs=interpolates,grad_outputs
                             =torch.ones_like(disc_interpolates),
                             create_graph=True,retain_graph=True)[0]
    gp=((gradients.norm(2,dim=1)-1)**2).mean()*LAMNDA #梯度惩罚变成常量
    return gp
def weights_init(m):
    if isinstance(m, nn.Linear):
        # m.weight.data.normal_(0.0, 0.02)
        nn.init.kaiming_normal_(m.weight)
        m.bias.data.fill_(0)
def main():
    torch.manual_seed(23)
    torch.cuda.manual_seed(23)#设置pytorch 随机化参数
    np.random.seed(23)#设置随机化种子
    device=torch.device('cuda')
    G=Generator().to(device)
    print(G)
    D=Discriminator().to(device)
    print(D)
    
    G.apply(weights_init)
    D.apply(weights_init)
    optim_G = optim.Adam(G.parameters(),lr=1e-3,betas=(0.5,0.9))
    optim_D = optim.Adam(D.parameters(), lr=1e-3, betas=(0.5, 0.9))
    data_iter=data_generator()
    print('batch',next(data_iter).shape)
    viz.line([[0,0]],[0],win='loss',opts=dict(title='loss',legend=['D','G']))
    
    for epoch in range(50000):
        #1. train discriminator for k steps
        #每个epoch 训练5次鉴别器 和1次生成器
        for _ in range(5):
            x=next(data_iter)
            xr=torch.from_numpy(x).cuda()#numpy数据转化为gpu tensor
            #[b]
            predr=(D(xr))
            #max log(lossr)
            lossr=-(predr.mean())#梯度下降法 取最小值
            
            #[b,2] 设计fake数据
            z=torch.randn(batchsz,2).cuda()
            #stop gradient on G
            #[b,2]
            xf =G(z).detach()
            #[b]
            predf=(D(xf))
            #min pedf 最小化数据预测
            lossf=(predf.mean())#loss 是标量
            
            #gradient penalty
            gp =gradient_penalty(D,xr,xf)
            
            loss_D= lossr+lossf+gp
            optim_D.zero_grad() #在训练时将D的梯度清零
            loss_D.backward()
            #for p in D.parameters()
            # print(p.grad.norm())
            optim_D.step()
        #2. train Generator    
        z=torch.randn(batchsz,2).cuda()
        xf=G(z)
        predf=(D(xf))
        
        # max predf
        loss_G=-(predf.mean())
        optim_G.zero_grad()
        loss_G.backward()
        optim_G.step()
        
        if epoch %100==0:
            viz.line([[loss_D.item(),loss_G.item()]],[epoch],win='loss',update='append')
            generator_image(D,G,xr,epoch)
            print(loss_D.item(),loss_G.item())
if __name__=='__main__':
    main()
        
        
        
            
            
    
    
    
    
    
    
       

GPU tensor不能直接转换成numpy先转换成CPU tensor，再转换成numpy

例如对GPU tensor img 操作转换成numpy为

img = img.cpu()

img = img.numpy()