1.在pytorch中训练GAN来生成数字.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "https://mp.weixin.qq.com/s?__biz=MzI5NTIxNTg0OA==&mid=2247498068&idx=3&sn=146f902866fa2094aca0d70989563ed3&chksm=ec544ed3db23c7c50bc370b93214219c9bd0354a58807de1763d1defb896988d6ec7d2bf9c9e&scene=0&xtrack=1&key=e8c75d77299604ef82eb6769e6b689f2ca2cec7932b55edf6f49266cca36c0541f2172f2b6a0f4c01cb62a62835ee855450c80a25548300ce516a1d8cf4d665c6952e90d059ee1d0988223de941c3a6a&ascene=1&uin=MjA1MjAyODkxNg%3D%3D&devicetype=Windows+10&version=62070141&lang=zh_CN&pass_ticket=Vsf3NZgYf5i5ZiliBW0AmsBXvt2ictUXr4fKbzPM5HHCZws0l68BCE4kvylUn040"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn,optim\n",
    "from torch.autograd.variable import Variable\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Preprocessing\n",
    "transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5,),(0.5,))])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training data\n",
    "train_set = torchvision.datasets.MNIST(root='../dataset/',train=True,download=True,transform=transform)\n",
    "train_loader = torch.utils.data.DataLoader(train_set,batch_size=32,shuffle=True)\n",
    "# Labels\n",
    "classes = [str(i) for i in range(0,10)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 判别器\n",
    "class Discriminator(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.model = nn.Sequential(\n",
    "            nn.Linear(784,1024),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Dropout(0.3),\n",
    "            nn.Linear(1024,512),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Dropout(0.3),\n",
    "            nn.Linear(512,256),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Dropout(0.3),\n",
    "            nn.Linear(256,1),\n",
    "            nn.Sigmoid()\n",
    "        )\n",
    "    def forward(self,x):\n",
    "        out = self.model(x.view(x.size(0),784))\n",
    "        out = out.view(out.size(0),-1)\n",
    "        return out.cuda()\n",
    "discriminator = Discriminator()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 生成器\n",
    "class Generator(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.model = nn.Sequential(\n",
    "            nn.Linear(100,256),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Linear(256,512),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Linear(512,1024),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Linear(1024,784),\n",
    "            nn.Tanh()\n",
    "        )\n",
    "    def forward(self,x):\n",
    "        x = x.view(x.size(0),100)\n",
    "        out = self.model(x).cuda()\n",
    "        return out\n",
    "\n",
    "generator = Generator()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using CUDA\n"
     ]
    }
   ],
   "source": [
    "if torch.cuda.is_available():\n",
    "    print('Using CUDA')\n",
    "    discriminator.cuda()\n",
    "    generator.cuda()\n",
    "# 设置损失函数和优化器\n",
    "lr = 0.0001\n",
    "num_epochs = 40\n",
    "num_batches = len(train_loader)\n",
    "\n",
    "criterion = nn.BCELoss()\n",
    "d_optimizer = torch.optim.Adam(discriminator.parameters(),lr=lr)\n",
    "g_optimizer = torch.optim.Adam(generator.parameters(),lr=lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 开始循环训练，pytorch中的训练循环通常由一个遍历各个训练周期的外部循环和一个遍历batch数据的内部循环组成，\n",
    "# 训练GAN的关键是我们需要在一个循环中更新生成器和判别器\n",
    "def train_discriminator(discriminator,real_images,real_labels,fake_images,fake_labels):\n",
    "    discriminator.zero_grad()\n",
    "    \n",
    "    # 真实图片的预测，损失和得分\n",
    "    predictions = discriminator(real_images)\n",
    "    real_loss = criterion(predictions,real_labels)\n",
    "    real_score = predictions\n",
    "    # 假的图片的预测，损失和得分\n",
    "    predictions = discriminator(fake_images)\n",
    "    fake_loss = criterion(predictions,fake_labels)\n",
    "    fake_score = predictions\n",
    "    # 计算总的损失，更新权重和优化器\n",
    "    d_loss = real_loss + fake_loss\n",
    "    d_loss.backward()\n",
    "    d_optimizer.step()\n",
    "    return d_loss,real_score,fake_score\n",
    "\n",
    "def train_generator(generator,discriminator_outputs,real_labels):\n",
    "    generator.zero_grad()\n",
    "    \n",
    "    # 计算总损失，更新权重和优化器\n",
    "    g_loss = criterion(discriminator_outputs,real_labels)\n",
    "    g_loss.backward()\n",
    "    g_optimizer.step()\n",
    "    return g_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 开始训练\n",
    "for epoch in range(num_epochs):\n",
    "    for n,(images,_) in enumerate(train_loader):\n",
    "        # (1)为判别器准备真实图片\n",
    "        real_images = Variable(images).cuda()\n",
    "        real_labels = Variable(torch.ones(images.size(0))).cuda()\n",
    "        # (2)为生成器准备随机噪声数据\n",
    "        noise = Variable(torch.randn(images.size(0),100)).cuda()\n",
    "        # (3)为判别器准备假的图片\n",
    "        fake_images = generator(noise)\n",
    "        fake_labels = Variable(torch.zeros(images.size(0))).cuda()\n",
    "        # (4)在真实图片和假的图片上训练判别器\n",
    "        d_loss,real_score,fake_score = train_discriminator(discriminator,real_images,real_labels,fake_images,fake_labels)\n",
    "        # (5a)从生成器产生一些新的假的图片\n",
    "        # (5b)假的图片在判别器上得到预测的标签\n",
    "        noise = Variable(torch.randn(images.size(0),100)).cuda()\n",
    "        fake_images = generator(noise)\n",
    "        \n",
    "        outputs = discriminator(fake_images)\n",
    "        # 训练生成器\n",
    "        g_loss = train_generator(generator,outputs,real_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:root] *",
   "language": "python",
   "name": "conda-root-py"
  },
  "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}