CGAN_mnist-02.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.autograd import Variable\n",
    "\n",
    "import torch.utils.data\n",
    "from torch import nn, optim\n",
    "from torchvision import datasets, transforms\n",
    "from torchvision.utils import save_image\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.preprocessing import LabelBinarizer\n",
    "import imageio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "project_root = os.path.realpath('.')\n",
    "print(project_root)\n",
    "os.chdir(project_root)\n",
    "\n",
    "no_cuda = False\n",
    "cuda_available = not no_cuda and torch.cuda.is_available()\n",
    "device = torch.device(\"cuda\" if cuda_available else \"cpu\")\n",
    "BATCH_SIZE = 64\n",
    "EPOCH = 100\n",
    "SEED = 8\n",
    "\n",
    "kwargs = {'num_workers': 1, 'pin_memory': True} if cuda_available else {}\n",
    "train_loader = torch.utils.data.DataLoader(\n",
    "    datasets.MNIST('./MNIST_data', train=True, download=True,\n",
    "                   transform=transforms.ToTensor()),\n",
    "    batch_size=BATCH_SIZE, shuffle=True, **kwargs)\n",
    "\n",
    "test_loader = torch.utils.data.DataLoader(\n",
    "    datasets.MNIST('./MNIST_data', train=False, transform=transforms.ToTensor()),\n",
    "    batch_size=BATCH_SIZE, shuffle=True, **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_classes = 10\n",
    "# Model params\n",
    "g_input_size = 100     # Random noise dimension coming into generator, per output vector\n",
    "\n",
    "g_output_size = 784    # size of generated output vector\n",
    "\n",
    "d_input_size = 784   # Minibatch size - cardinality of distributions\n",
    "\n",
    "d_output_size = 1    # Single dimension for 'real' vs. 'fake'\n",
    "\n",
    "\n",
    "d_learning_rate = 2e-4  # 2e-4\n",
    "g_learning_rate = 2e-4\n",
    "optim_betas = (0.9, 0.999)\n",
    "\n",
    "print_interval = 200\n",
    "\n",
    "d_steps = 1  # 'k' steps in the original GAN paper. Can put the discriminator on higher training freq than generator\n",
    "g_steps = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Generator(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.fc1_1 = nn.Linear(100, 256)\n",
    "        self.fc1_1_bn = nn.BatchNorm1d(256)\n",
    "        self.fc1_2 = nn.Linear(10, 256)\n",
    "        self.fc1_2_bn = nn.BatchNorm1d(256)\n",
    "        self.fc2 = nn.Linear(512, 512)\n",
    "        self.fc2_bn = nn.BatchNorm1d(512)\n",
    "        self.fc3 = nn.Linear(512, 1024)\n",
    "        self.fc3_bn = nn.BatchNorm1d(1024)\n",
    "        self.fc4 = nn.Linear(1024, 784)\n",
    "\n",
    "\n",
    "        for m in self.modules():\n",
    "            if isinstance(m, nn.Linear):\n",
    "                nn.init.xavier_normal_(m.weight)\n",
    "            elif isinstance(m, nn.BatchNorm2d):\n",
    "                nn.init.constant_(m.weight, 1)\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "                \n",
    "    def forward(self, input, label):\n",
    "        x = F.relu(self.fc1_1_bn(self.fc1_1(input)))\n",
    "        y = F.relu(self.fc1_2_bn(self.fc1_2(label)))\n",
    "        x = torch.cat([x, y], 1)\n",
    "        x = F.relu(self.fc2_bn(self.fc2(x)))\n",
    "        x = F.relu(self.fc3_bn(self.fc3(x)))\n",
    "        x = F.tanh(self.fc4(x))\n",
    "        \n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Discriminator(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.fc1_1 = nn.Linear(784, 1024)\n",
    "        self.fc1_2 = nn.Linear(10, 1024)\n",
    "        self.fc2 = nn.Linear(2048, 512)\n",
    "        self.fc2_bn = nn.BatchNorm1d(512)\n",
    "        self.fc3 = nn.Linear(512, 256)\n",
    "        self.fc3_bn = nn.BatchNorm1d(256)\n",
    "        self.fc4 = nn.Linear(256, 1)\n",
    "        \n",
    "        for m in self.modules():\n",
    "            if isinstance(m, nn.Linear):\n",
    "                nn.init.xavier_normal_(m.weight)\n",
    "            elif isinstance(m, nn.BatchNorm2d):\n",
    "                nn.init.constant_(m.weight, 1)\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "                \n",
    "    def forward(self, input, label):\n",
    "        x = F.leaky_relu(self.fc1_1(input), 0.2)\n",
    "        y = F.leaky_relu(self.fc1_2(label), 0.2)\n",
    "        x = torch.cat([x, y], 1)\n",
    "        x = F.leaky_relu(self.fc2_bn(self.fc2(x)), 0.2)\n",
    "        x = F.leaky_relu(self.fc3_bn(self.fc3(x)), 0.2)\n",
    "        x = F.sigmoid(self.fc4(x))\n",
    "        \n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "D  = Discriminator().to(device)\n",
    "G = Generator().to(device)\n",
    "print(D)\n",
    "print(G)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = nn.BCELoss()\n",
    "lr = 0.0002\n",
    "d_optimizer = torch.optim.Adam(D.parameters(), lr=lr, betas=(0.5, 0.999))\n",
    "g_optimizer = torch.optim.Adam(G.parameters(), lr=lr, betas=(0.5, 0.999))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mini_batch = 5\n",
    "y_=torch.LongTensor([3,2,4,1,5])\n",
    "y_real_ = torch.ones(mini_batch)\n",
    "\n",
    "y_fake_ = torch.zeros(mini_batch)\n",
    "\n",
    "y_label_ = torch.zeros(mini_batch, 10)\n",
    "y_label_.scatter_(1, y_.view(mini_batch, 1), 1)\n",
    "y_label_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lb = LabelBinarizer()\n",
    "lb.fit(list(range(0,n_classes)))\n",
    "   #将标签进行one-hot编码\n",
    "def to_categrical(y: torch.FloatTensor):\n",
    "    y_n = y.numpy()\n",
    "    \n",
    "    y_one_hot = lb.transform(y_n)\n",
    "    floatTensor = torch.FloatTensor(y_one_hot)\n",
    "    return floatTensor\n",
    "\n",
    "#样本和one-hot标签进行连接，以此作为条件生成\n",
    "def concanate_data_label(x, y):\n",
    "#     print(f'dimenion y: {y.shape} {x.shape}')\n",
    "    y_one_hot = to_categrical(y)\n",
    "#     print(y_one_hot.shape)\n",
    "    con = torch.cat((x, y_one_hot), 1)\n",
    "#     print(con.shape)\n",
    "    \n",
    "    return con.to(device)\n",
    "\n",
    "D_losses = []\n",
    "G_losses = []\n",
    "\n",
    "def train(epoch):\n",
    "    D.train()\n",
    "    G.train()\n",
    "    \n",
    "     # learning rate decay\n",
    "    if (epoch+1) == 30:\n",
    "        g_optimizer.param_groups[0]['lr'] /= 10\n",
    "        d_optimizer.param_groups[0]['lr'] /= 10\n",
    "        print(\"learning rate change!\")\n",
    "\n",
    "    if (epoch+1) == 40:\n",
    "        g_optimizer.param_groups[0]['lr'] /= 10\n",
    "        d_optimizer.param_groups[0]['lr'] /= 10\n",
    "        print(\"learning rate change!\")\n",
    "    \n",
    "    D_losses.clear()\n",
    "    G_losses.clear()\n",
    "    for batch_idx, (data, label) in enumerate(train_loader):\n",
    "        \n",
    "        D.zero_grad()\n",
    "        data_flatten = data.view(data.size(0), d_input_size)\n",
    "        label_categorical = to_categrical(label)\n",
    "        \n",
    "        for d_index in range(d_steps):\n",
    "            # 1. Train D on real+fake\n",
    "            #  1A: Train D on real\n",
    "            \n",
    "            d_real_decision = D(Variable(data_flatten), label_categorical)\n",
    "           \n",
    "            d_real_error = criterion(d_real_decision, Variable(torch.ones(data.shape[0],1)).to(device))  # ones = true\n",
    "            d_real_error.backward() # compute/store gradients, but don't change params\n",
    "            \n",
    "            #  1B: Train D on fake\n",
    "            d_gen_input = Variable(torch.randn(data.shape[0], g_input_size))\n",
    "            d_fake_data = G(d_gen_input.to(device), label_categorical).detach()  # detach to avoid training G on these labels\n",
    "            \n",
    "            d_fake_decision = D(d_fake_data.to(device), label_categorical)\n",
    "            d_fake_error = criterion(d_fake_decision, Variable(torch.zeros(data.shape[0], 1)).to(device))  # zeros = fake\n",
    "            d_fake_error.backward()\n",
    "            d_optimizer.step()     # Only optimizes D's parameters; changes based on stored gradients from backward()\n",
    "            \n",
    "            d_error = d_real_error + d_fake_error\n",
    "            D_losses.append(d_error.data)\n",
    "\n",
    "        for g_index in range(g_steps):\n",
    "        # 2. Train G on D's response (but DO NOT train D on these labels)\n",
    "            G.zero_grad()\n",
    "\n",
    "            gen_input = Variable(torch.randn(data.shape[0], g_input_size).to(device))\n",
    "            g_fake_data = G(gen_input, label_categorical)\n",
    "            dg_fake_decision = D(g_fake_data.to(device), label_categorical)\n",
    "            g_error = criterion(dg_fake_decision, Variable(torch.ones(data.shape[0], 1)).to(device))  # we want to fool, so pretend it's all genuine\n",
    "\n",
    "            g_error.backward()\n",
    "            g_optimizer.step()  # Only optimizes G's parameters\n",
    "            \n",
    "            G_losses.append(g_error.data)\n",
    "        \n",
    "            \n",
    "    print('[%d/%d]: D(x): %.3f, D(G(z)): %.3f' % (\n",
    "        epoch , EPOCH,np.mean(D_losses), np.mean(G_losses)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(epoch):\n",
    "    G.eval()\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_samples():\n",
    "    '''10x10 matrix,row 0-9 represent label 0-9'''\n",
    "    sample = torch.randn(100, g_input_size) #100x100  \n",
    "    label = torch.zeros(100, 10)\n",
    "    \n",
    "    for i in range(100):\n",
    "        label[i, i % n_classes] = i // n_classes\n",
    "    \n",
    "    return sample.to(device), label.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "11//10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "des_path = os.path.join(project_root, 'cgan_new_results/')\n",
    "if not os.path.exists(des_path):\n",
    "    os.makedirs(des_path, exist_ok=True)\n",
    "\n",
    "import math,  itertools\n",
    "from IPython import display\n",
    "#for plot\n",
    "size_figure_grid = 10\n",
    "fig, ax = plt.subplots(size_figure_grid, size_figure_grid, figsize=(6, 6))\n",
    "for i, j in itertools.product(range(size_figure_grid), range(size_figure_grid)):\n",
    "    ax[i,j].get_xaxis().set_visible(False)\n",
    "    ax[i,j].get_yaxis().set_visible(False)\n",
    "\n",
    "#train\n",
    "for epoch in range(1, EPOCH + 1):\n",
    "    display.clear_output(wait=True)\n",
    "    train(epoch)\n",
    "#     test(epoch)\n",
    "    with torch.no_grad():\n",
    "        sample, label = get_samples()\n",
    "        sample = G(sample, label).cpu()\n",
    "        save_image(sample.view(100, 1, 28, 28),\n",
    "                   f'{des_path}epoch_{epoch}.png', nrow=10)\n",
    "        \n",
    "        for k in range(100):\n",
    "            i = k//10\n",
    "            j = k%10\n",
    "            ax[i,j].cla()\n",
    "            ax[i,j].imshow(sample[k,:].data.cpu().numpy().reshape(28, 28),cmap='Greys')\n",
    "        \n",
    "        display.display(plt.gcf())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "images = []\n",
    "for i in range(1, EPOCH+1):\n",
    "    one_image = f'{des_path}epoch_{i}.png'\n",
    "    images.append(imageio.imread(one_image))\n",
    "imageio.mimsave(f'{des_path}cgan_new.gif', images, fps=5)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}