dcgan_torch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 读取数据集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.datasets import CIFAR10 #下载图片数据集\n",
    "from torch.utils.data import DataLoader #读取批次\n",
    "import torchvision.transforms as transforms #张量转换\n",
    "from torchvision.utils import save_image #保存图片"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#读取数据集\n",
    "dataset = CIFAR10(root = './data', \n",
    "                 download = True, transform = transforms.ToTensor()) #下载数据集\n",
    "dataloader = DataLoader(dataset, batch_size= 64, shuffle= True) #按批次读取数据(一批64张，总共有50000张，所以有50000/64=781+1=782批)，shuffle= True打乱数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for batch_idx, data in enumerate(dataloader):\n",
    "    real_images, _ = data\n",
    "    batch_size = real_images.shape[0]\n",
    "    print('#{} has {} images.'.format(batch_idx, batch_size))\n",
    "    if batch_idx % 100 == 0:\n",
    "        path = './data/CIFAR10_shuffled_batch{:03d}.png'.format(batch_idx)\n",
    "        save_image(real_images, path, nrow=8, normalize=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#对上面for循环的解释：\n",
    "#batch_idx 是第几批，batch_size是一批有几个（64个，最后一批凑不齐只有16个）也可以用real_images.size(0)来表示。\n",
    "#save_image是保存图片，使用if条件语句，每100批保存当批的图片，共有782批，所以会保存8张图。每一张图有64张小图，设置显示行数为8，就会有8*8=64张小图凑成的一张大图，normallize设置是否调整输张量的范围到[0,1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 设置网络"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#构建鉴别网络\n",
    "n_d_feature = 64 #潜在大小\n",
    "n_channel = 3 #输入通道数\n",
    "dnet = nn.Sequential(\n",
    "        nn.Conv2d(n_channel, n_d_feature, kernel_size=4,\n",
    "                 stride=2, padding=1),\n",
    "        nn.LeakyReLU(0.2),\n",
    "        nn.Conv2d(n_d_feature, 2 * n_d_feature, kernel_size=4,\n",
    "                 stride=2, padding=1, bias=False),\n",
    "        nn.BatchNorm2d(2 * n_d_feature),\n",
    "        nn.LeakyReLU(0.2),\n",
    "        nn.Conv2d(2 * n_d_feature, 4 * n_d_feature, kernel_size=4,\n",
    "                 stride=2, padding=1, bias=False),\n",
    "        nn.BatchNorm2d(4 * n_d_feature),\n",
    "        nn.LeakyReLU(0.2),\n",
    "        nn.Conv2d(4 * n_d_feature, 1, kernel_size=4))\n",
    "print(dnet)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#构建生成网络\n",
    "latent_size = 64 #输入通道数\n",
    "n_g_feature = 64 #输出通道数\n",
    "\n",
    "gnet = nn.Sequential(\n",
    "        nn.ConvTranspose2d(latent_size, 4 * n_g_feature, kernel_size=4,\n",
    "                          bias=False),\n",
    "        nn.BatchNorm2d(4 * n_g_feature),\n",
    "        nn.ReLU(),\n",
    "        nn.ConvTranspose2d(4 * n_g_feature, 2 * n_g_feature, kernel_size=4,\n",
    "                          stride=2, padding=1, bias=False),\n",
    "        nn.BatchNorm2d(2 * n_g_feature),\n",
    "        nn.ReLU(),\n",
    "        nn.ConvTranspose2d(2 * n_g_feature, n_g_feature, kernel_size=4,\n",
    "                          stride=2, padding=1, bias=False),\n",
    "        nn.BatchNorm2d(n_g_feature),\n",
    "        nn.ReLU(),\n",
    "        nn.ConvTranspose2d(n_g_feature, n_channel, kernel_size=4,\n",
    "                          stride=2, padding=1),\n",
    "        nn.Sigmoid())\n",
    "print(gnet)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#网络初始化\n",
    "import torch.nn.init as init\n",
    "def weights_init(m):\n",
    "    if type(m) in [nn.ConvTranspose2d, nn.Conv2d]:\n",
    "        init.xavier_normal_(m.weight)\n",
    "    elif type(m) == nn.BatchNorm2d:\n",
    "        init.normal_(m.weight, 1.0, 0.02)\n",
    "        init.constant_(m.bias, 0)\n",
    "        \n",
    "gnet.apply(weights_init)\n",
    "dnet.apply(weights_init)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 训练网络"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.optim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#载入cpu训练的预参数\n",
    "#checkpoint_d = torch.load('discriminator.pth', map_location='cpu')\n",
    "#checkpoint_g = torch.load('generator.pth', map_location='cpu')\n",
    "#dnet.load_state_dict(checkpoint_d)\n",
    "#gnet.load_state_dict(checkpoint_g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#定义损失\n",
    "criterion = nn.BCEWithLogitsLoss()\n",
    "#定义优化器\n",
    "goptimizer = torch.optim.Adam(gnet.parameters(),\n",
    "                             lr=0.0002, betas=(0.5, 0.999))\n",
    "doptimizer = torch.optim.Adam(dnet.parameters(),\n",
    "                             lr=0.0002, betas=(0.5, 0.999))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#生成噪音数据，输入到G网络的数据\n",
    "batch_size = 64\n",
    "fixed_noises = torch.randn(batch_size, latent_size, 1, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#开始训练\n",
    "epoch_num = 5#共训练10个周期\n",
    "for epoch in range(epoch_num):\n",
    "    for batch_idx, data in enumerate(dataloader):\n",
    "        real_images, _ = data\n",
    "        batch_size = real_images.shape[0]\n",
    "        #训练判别器D\n",
    "        labels = torch.ones(batch_size)#真实数据的标签：1\n",
    "        preds = dnet(real_images) #将真实数据喂给D网络\n",
    "        outputs = preds.reshape(-1) #转换成未知行\n",
    "        dloss_real = criterion(outputs, labels)\n",
    "        dmean_real = outputs.sigmoid().mean() #计算判别器将多少真数据判别为真，仅用于输出显示\n",
    "        \n",
    "        noises = torch.randn(batch_size, latent_size, 1, 1)\n",
    "        fake_images = gnet(noises) #生成假数据\n",
    "        labels = torch.zeros(batch_size)#生成假数据的标签：0\n",
    "        fake = fake_images.detach() #类似于固定生成器参数\n",
    "        preds = dnet(fake) #将假数据喂给判别器\n",
    "        outputs = preds.reshape(-1)#转换成未知行\n",
    "        dloss_fake = criterion(outputs, labels)\n",
    "        dmean_fake = outputs.sigmoid().mean() #计算判别器将多少假数据判断为真，仅用于输出显示\n",
    "        \n",
    "        dloss = dloss_real + dloss_fake #总的鉴别器损失为两者之和\n",
    "        dnet.zero_grad()#梯度清零\n",
    "        dloss.backward()#反向传播\n",
    "        doptimizer.step()\n",
    "        \n",
    "        #训练生成器G\n",
    "        labels = torch.ones(batch_size)#在训练生成器G时，希望生成器的标签为1\n",
    "        preds = dnet(fake_images)#让假数据通过鉴别网络\n",
    "        outputs = preds.reshape(-1)#转换成未知行\n",
    "        gloss = criterion(outputs, labels)\n",
    "        gmean_fake = outputs.sigmoid().mean() #计算判别器将多少假数据判断为真，仅用于输出显示\n",
    "        \n",
    "        gnet.zero_grad()#梯度清零\n",
    "        gloss.backward()#反向传播\n",
    "        goptimizer.step()\n",
    "        \n",
    "        #输出本步训练结果\n",
    "        print('[{}/{}]'.format(epoch, epoch_num) + '[{}/{}]'.format(batch_idx, len(dataloader)) +\n",
    "             '鉴别器G损失:{:g} 生成器D损失：{:g}'.format(dloss, gloss) + \n",
    "             '真数据判真比例：{:g} 假数据判真比例：{:g}/{:g}'.format(dmean_real, dmean_fake, gmean_fake))\n",
    "        if batch_idx % 100 == 0:\n",
    "            fake = gnet(fixed_noises) #噪声生成假数据\n",
    "            path = './data/images_epoch{:02d}_batch{:03d}.png'.format(epoch, batch_idx)\n",
    "            save_image(fake, path, normalize=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 保存模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(gnet.state_dict(),'./generator.pth')\n",
    "torch.save(dnet.state_dict(),'./discriminator.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}