【轩说Pytorch】朴素GAN 代码实战（2gpu数据并行＋tqdm可视化）

GAN代码实操

GAN的理论部分已经讲过，下面是代码实战。可以生成MNIST数据集的手写数字。

1. 导包

   • torch是pytorch的主目录包，用于选取优化器torch.optim.Adam,torch.optim.SGD，选取模型torch.nn，建立张量torch.Tensor,torch.ones,torch.randn，选取设备torch.device,torch.cuda，构建数据torch.utils.data，保存数据torch.utils.save_image等

   • nn是neural network的缩写，建模使用。比如nn.Module,nn.Sequence等，以及相应的损失函数nn.CrossEntropyLoss,nn.BCELoss等。

   • torchvision是pytorch的一个图形库，它服务于PyTorch深度学习框架的，主要用来构建计算机视觉模型。torchvision.datasets: 一些加载数据的函数及常用的数据集接口；torchvision.models: 包含常用的模型结构（含预训练模型），例如AlexNet、VGG、ResNet等；torchvision.transforms: 常用的图片变换，例如裁剪、旋转等；torchvision.utils: 其他的一些有用的方法。

   • argparse库是用来给python脚本传入参数的库

   • os库提供通用的、基本的操作系统交互功能

   • torchsummary库可以查看模型结构

   • tqdm库可以动态展示训练过程

   import torch
   from torch import nn
   import torchvision
   from torchvision import transforms
   from torch.utils import data
   import argparse
   import  os
   from torchsummary import summary
   import tqdm
   

2. 导入参数＋建立输出文件夹

   parser=argparse.ArgumentParser()
   parser.add_argument("--batchsize",type=int,default=128,help="input the size of batch")
   parser.add_argument("--epochs",type=int,default=200,help="input the num of epoches")
   parser.add_argument("--outdir",default="samples_复习",help="input the dir path of samples")
   args=parser.parse_args()
   hidden_size=256
   latent_size=64
   image_size=784
   if not os.path.isdir(args.outdir):
       os.mkdir(args.outdir)
   
   

3. 设定cuda

   详见(32条消息) 【轩说Pytorch】用GPU训练模型_留下一些记忆的博客-CSDN博客

   os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
   os.environ["CUDA_VISIBLE_DEVICES"] = "0,1" #指定只可见0号物理卡和1号物理卡，可省略
   device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")#指定0号逻辑卡为device
   

4. 加载数据

   由于G的最后一层为$nn.Tanh()$，所以生成的图像的像素值分布在[-1,1]区间，为得到相同范围的分布，将$P_{data}(X)$的$X$也映射为[-1,1]的取值，故采用$transforms.Normalize(mean=0.5,std=0.5)$。

   $X_{Image}(0,255)$->$X_{Tensor}(0,1)$->$X_{normalized}(-1,1)$

   数据集采用MNSIT手写数字，并封装在dataloader中。

   transform=transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=0.5,std=0.5)])
   dataset=torchvision.datasets.MNIST(root="../data/",download=True,train=True,transform=transform)
   dataloader=data.DataLoader(dataset=dataset,batch_size=args.batchsize,shuffle=False)
   
   print("\nfeature size:",next(iter(dataloader))[0][0].shape)
   """
   feature size: torch.Size([1, 28, 28])
   """
   

5. 定义模型并将模型移到device上

   最朴素的GAN。G:64-256-256-784，D:784-256-256-1

   G=nn.Sequential(nn.Linear(latent_size,hidden_size),nn.ReLU(),
                  nn.Linear(hidden_size,hidden_size),nn.ReLU(),
                  nn.Linear(hidden_size,image_size),nn.Tanh())
   
   D=nn.Sequential(nn.Linear(image_size,hidden_size),nn.ReLU(),
                   nn.Linear(hidden_size,hidden_size),nn.ReLU(),
                   nn.Linear(hidden_size,1),nn.Sigmoid())
   D=D.to(device)
   G=G.to(device)
   

6. 做DataParallel数据并行

   device_ids=[0,1]
   D=nn.DataParallel(module=D,device_ids=device_ids)
   G=nn.DataParallel(module=G,device_ids=device_ids)
   

7. 定义损失函数和优化器

   D_optim=torch.optim.Adam(D.parameters(),lr=0.0002)
   G_optim=torch.optim.Adam(G.parameters(),lr=0.0002)
   criterion=nn.BCELoss()#二元交叉熵损失函数
   batch_size=args.batchsize
   num_epochs=args.epochs
   k=1
   

8. 反归一化

   把[-1,1]的Tensor转为[0,1]的Tensor。clamp(*)函数的功能：将输入input张量每个元素的夹紧到区间$[min,max]$并返回结果到一个新张量。

   def denorm(X):
       return ((X+1)/2).clamp(0,1)
   

9. 开始训练

   (32条消息) tqdm的使用和例子_Le0v1n的博客-CSDN博客

   使用tqdm作为可视化训练进度的库。定义$progressbar=tqdm.tqdm(如下参数)$

   iterable: 可迭代的对象, 在⼿动更新时不需要进⾏设置
   desc: 字符串, 左边进度条描述⽂字
   total: 总的项⽬数（一般不写，后台自动计算，写的话容易出错）
   leave: bool值, 迭代完成后是否保留进度条（默认False）
   file: 输出指向位置, 默认是终端, ⼀般不需要设置
   ncols: 调整进度条宽度, 默认是根据环境⾃动调节长度, 如果设置为0, 就没有进度条, 只有输出的信息
   unit: 描述处理项⽬的⽂字, 默认是it, 例如: 100 it/s, 处理照⽚的话设置为img ,则为 100 img/s
   unit_scale: ⾃动根据国际标准进⾏项⽬处理速度单位的换算, 例如 100000 it/s >> 100k it/s
   colour: 进度条颜色
   

   在训练过程中动态自定义进度条显示信息

   1. set_description()
   2. set_postfix()

   set_postfix形参是用的kwargs传参，所以我们可以：

   • 用字典传参 -> pbar.set_postfix({“key_1”: “value_1”, …})
   • 直接用关键字传参 -> pbar.set_postfix(key_1 = value_1, key_2 = value_2, …)
   • 混着用 -> pbar.set_postfix({“key_1”: value_1, “key_2”: value_2, …}, key_3 = value_3, …)

   每过10个epoch，就将生成的图片保存一下，方便观察效果。深度学习模型中，一般使用torchvision.utils中的save_image(tensor, PATH)函数进行图像保存，这种方式只能保存RGB彩色图像，如果网络的输出是单通道灰度图像，则该函数依然会输出三个通道，每个通道的数值都是相同的，即“伪灰度图像”，虽然从视觉效果上看不出区别，但是图像所占内存比正常情况大了两倍。

   for epoch in range(num_epochs):
       pbar=tqdm.tqdm(iterable=dataloader,desc=f"Epoch {epoch+1}")
       for X,_ in pbar:
   		"""
   			中间部分见下一段
   		"""
           pbar.set_postfix({"G_loss":G_loss.item(),"D_loss":D_loss.item(),
                             "avg_D(x)":y_real.mean().item(),"avg_D(G(z))":y_gen.mean().item()})
           #.item()把Tensor(size=(1))->Scalar（标量）
   
       if epoch%10==0:
           X_gen=X_gen.reshape(X_gen.size(0),1,28,28)
           torchvision.utils.save_image(denorm(X_gen),os.path.join(args.outdir,'fake_images-{}.png'.format(epoch+1)))
   

   

   网络优化（for循环内部）

   • 训练D

           real_label=torch.ones(size=(batch_size,1)).to(device)
           fake_label=torch.zeros(size=(batch_size,1)).to(device)
           X=X.reshape((X.size(0),-1)).to(device)
           y_real=D(X)
           loss_real=criterion(y_real,real_label[:y_real.size(0)])
           z=torch.randn(size=(batch_size,latent_size)).to(device)
           X_gen=G(z)
           y_gen=D(X_gen)
           loss_fake=criterion(y_gen,fake_label[:y_gen.size(0)])
           D_loss=loss_fake+loss_real
           D_optim.zero_grad()
           D_loss.backward()
           D_optim.step()
   

   • 训练G

           z = torch.randn(size=(batch_size, latent_size)).to(device)
           X_gen = G(z)
           y_gen = D(X_gen)
           G_loss = criterion(y_gen, real_label[:y_gen.size(0)])
           G_optim.zero_grad()
           G_loss.backward()
           G_optim.step()
   

训练成果展示

第170个epoch过后：

数据并行化

程序/envs/pt/bin/python3.8是我的进程，在并行开启后，会出现两个进程，且2卡的memory占用量明显增加。

用单卡运行

用多卡做数据并行DataParallel