Main.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Mac-Torch-Gpu-VAE\n",
    "Use MPS (provided in [Pytorch v1.12.0+](https://pytorch.org/blog/introducing-accelerated-pytorch-training-on-mac/)) to accelerate the calculations in ARM Macbook.\n",
    "\n",
    "## 1. Define Parameters\n",
    "Parameters for training the VAE."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "\n",
    "parser = argparse.ArgumentParser(description='MAC Torch GPU VAE Example')\n",
    "parser.add_argument('--batch-size', type=int, default=128, metavar='N',\n",
    "                    help='input batch size for training (default: 128)')\n",
    "parser.add_argument('--epochs', type=int, default=3, metavar='N',\n",
    "                    help='number of epochs to train (default: 10)')\n",
    "parser.add_argument('--no-GPU', action='store_true', default=False,\n",
    "                    help='disables GPU training') # Use CPU or GPU? True:False;\n",
    "parser.add_argument('--seed', type=int, default=1, metavar='S',\n",
    "                    help='random seed (default: 1)')\n",
    "parser.add_argument('--log-interval', type=int, default=10, metavar='N',\n",
    "                    help='how many batches to wait before logging training status')\n",
    "args = parser.parse_args(args=[])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Parameters for using GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "args.GPU:True\n",
      "kwargs:{'num_workers': 2, 'pin_memory': True}\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "# Mac: mps\n",
    "args.GPU = not args.no_GPU and torch.backends.mps.is_available()\n",
    "# Linux: cuda\n",
    "# args.GPU = not args.no_GPU and torch.cuda.is_available()\n",
    "\n",
    "print(\"args.GPU:\" + str(args.GPU))\n",
    "\n",
    "torch.manual_seed(args.seed) # Fix inital parameters of model\n",
    "# Mac: mps\n",
    "device = torch.device(\"mps\" if args.GPU else \"cpu\") # Choose CPU or GPU device\n",
    "# Linux: cuda\n",
    "# device = torch.device(\"cuda\" if args.GPU else \"cpu\")\n",
    "\n",
    "kwargs = {\"num_workers\":2, \"pin_memory\":True} if args.GPU else {}\n",
    "print(\"kwargs:\" + str(kwargs))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. VAE Model Development"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import nn, optim\n",
    "from torch.nn import functional as F\n",
    "\n",
    "# VAE Model\n",
    "class VAE(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(VAE, self).__init__()\n",
    "\n",
    "        self.fc1 = nn.Linear(28*28, 400)\n",
    "        self.fc21 = nn.Linear(400, 20) # Mu\n",
    "        self.fc22 = nn.Linear(400, 20) # log(Var)\n",
    "        self.fc3 = nn.Linear(20, 400)\n",
    "        self.fc4 = nn.Linear(400, 28*28)\n",
    "    \n",
    "    def encoder(self, x):\n",
    "        h1 = F.relu(self.fc1(x))\n",
    "        mu = self.fc21(h1)\n",
    "        logvar = self.fc22(h1)\n",
    "        return mu, logvar\n",
    "    \n",
    "    def reparameterize(self, mu, logvar):\n",
    "        std = torch.exp(0.5*logvar)\n",
    "        eps = torch.randn_like(std)\n",
    "        return mu + eps*std\n",
    "    \n",
    "    def decode(self, z):\n",
    "        h3 = F.relu(self.fc3(z))\n",
    "        return torch.sigmoid(self.fc4(h3))\n",
    "    \n",
    "    def forward(self, x):\n",
    "        mu, logvar = self.encoder(x.view(-1, 28*28))\n",
    "        z = self.reparameterize(mu, logvar)\n",
    "        return self.decode(z), mu, logvar\n",
    "\n",
    "model = VAE().to(device)\n",
    "optimizer = optim.Adam(model.parameters(), lr = 1e-4) # Adam optimizer\n",
    "\n",
    "# Loss_function = -ELBO = \n",
    "def loss_function(x_hat, x, mu, logvar):\n",
    "    # Reconstruction Term\n",
    "    Recon = F.binary_cross_entropy(x_hat, x.view(-1, 28*28), reduction='sum')\n",
    "\n",
    "    # Regularization Loss \n",
    "    KL = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())\n",
    "\n",
    "    return Recon + KL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Train and Test Processes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.utils.data\n",
    "from torchvision import datasets, transforms\n",
    "from torchvision.utils import save_image\n",
    "\n",
    "# Dataloader for training and testing (read img as chw format)\n",
    "train_loader = torch.utils.data.DataLoader(\n",
    "    datasets.MNIST('./data', train = True, download = False, transform = transforms.ToTensor()),\n",
    "    batch_size = args.batch_size, shuffle = True, **kwargs)\n",
    "\n",
    "test_loader = torch.utils.data.DataLoader(\n",
    "    datasets.MNIST('./data', train = False, transform = transforms.ToTensor()),\n",
    "    batch_size = args.batch_size, shuffle = False, **kwargs)\n",
    "\n",
    "def train(epoch):\n",
    "    model.train() # Enable: batch normalization + drop out\n",
    "    train_loss = 0 # total loss in an epoch\n",
    "    for batch_idx, (data, _) in enumerate(train_loader):\n",
    "        # Initialization\n",
    "        data = data.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        # Loss backward\n",
    "        x_hat, mu, logvar = model(data)\n",
    "        loss = loss_function(x_hat, data, mu, logvar)\n",
    "        loss.backward()\n",
    "        train_loss += loss.item()\n",
    "        # Optimize\n",
    "        optimizer.step()\n",
    "        # Print Information (Batch)\n",
    "        if batch_idx % args.log_interval == 0:\n",
    "            print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss:{:.6f}'.format(\n",
    "                epoch, batch_idx * len(data), len(train_loader.dataset),\n",
    "                100. * batch_idx/len(train_loader),\n",
    "                loss.item() / len(data)))\n",
    "    \n",
    "    # Print Information (Epoch)\n",
    "    print(\"[Epoch]:{}, Average Loss:{:.4f}\".format(\n",
    "        epoch, train_loss / len(train_loader.dataset)))\n",
    "\n",
    "def test(epoch):\n",
    "    model.eval() # Disable: batch normalization + drop out\n",
    "    test_loss = 0 # total loss in an epoch\n",
    "    with torch.no_grad():\n",
    "        for i, (data, _) in enumerate(test_loader):\n",
    "            # Initialization\n",
    "            data = data.to(device)\n",
    "            # Loss Calculation\n",
    "            x_hat, mu, logvar = model(data)\n",
    "            test_loss += loss_function(x_hat, data, mu, logvar).item()\n",
    "            # Print Information (The First Batch)\n",
    "            if i == 0:\n",
    "                n = min(data.size(0), 8) # n samples shown\n",
    "                # Compare x and x_hat\n",
    "                comparison = torch.cat([data[:n],\n",
    "                    x_hat.view(args.batch_size, 1, 28, 28)[:n]])\n",
    "                save_image(comparison.cpu(), './results/rec_' + str(epoch) + '.png', nrows=n)\n",
    "    \n",
    "    test_loss /= len(test_loader.dataset)\n",
    "    print(\"Test Loss:{:.6f}\".format(test_loss))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Main()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Epoch: 1 [0/60000 (0%)]\tLoss:550.186035\n",
      "Train Epoch: 1 [1280/60000 (2%)]\tLoss:523.721436\n",
      "Train Epoch: 1 [2560/60000 (4%)]\tLoss:497.339233\n",
      "Train Epoch: 1 [3840/60000 (6%)]\tLoss:463.267517\n",
      "Train Epoch: 1 [5120/60000 (9%)]\tLoss:424.521362\n",
      "Train Epoch: 1 [6400/60000 (11%)]\tLoss:374.330780\n",
      "Train Epoch: 1 [7680/60000 (13%)]\tLoss:333.500946\n",
      "Train Epoch: 1 [8960/60000 (15%)]\tLoss:301.895233\n",
      "Train Epoch: 1 [10240/60000 (17%)]\tLoss:291.794464\n",
      "Train Epoch: 1 [11520/60000 (19%)]\tLoss:282.959015\n",
      "Train Epoch: 1 [12800/60000 (21%)]\tLoss:269.283905\n",
      "Train Epoch: 1 [14080/60000 (23%)]\tLoss:262.155670\n",
      "Train Epoch: 1 [15360/60000 (26%)]\tLoss:268.469421\n",
      "Train Epoch: 1 [16640/60000 (28%)]\tLoss:257.672729\n",
      "Train Epoch: 1 [17920/60000 (30%)]\tLoss:248.334396\n",
      "Train Epoch: 1 [19200/60000 (32%)]\tLoss:247.576370\n",
      "Train Epoch: 1 [20480/60000 (34%)]\tLoss:242.984421\n",
      "Train Epoch: 1 [21760/60000 (36%)]\tLoss:233.088470\n",
      "Train Epoch: 1 [23040/60000 (38%)]\tLoss:243.873383\n",
      "Train Epoch: 1 [24320/60000 (41%)]\tLoss:238.423233\n",
      "Train Epoch: 1 [25600/60000 (43%)]\tLoss:238.230072\n",
      "Train Epoch: 1 [26880/60000 (45%)]\tLoss:231.435211\n",
      "Train Epoch: 1 [28160/60000 (47%)]\tLoss:234.785843\n",
      "Train Epoch: 1 [29440/60000 (49%)]\tLoss:223.958221\n",
      "Train Epoch: 1 [30720/60000 (51%)]\tLoss:222.026489\n",
      "Train Epoch: 1 [32000/60000 (53%)]\tLoss:224.456284\n",
      "Train Epoch: 1 [33280/60000 (55%)]\tLoss:229.785797\n",
      "Train Epoch: 1 [34560/60000 (58%)]\tLoss:221.096588\n",
      "Train Epoch: 1 [35840/60000 (60%)]\tLoss:217.251984\n",
      "Train Epoch: 1 [37120/60000 (62%)]\tLoss:223.627380\n",
      "Train Epoch: 1 [38400/60000 (64%)]\tLoss:214.305313\n",
      "Train Epoch: 1 [39680/60000 (66%)]\tLoss:219.063477\n",
      "Train Epoch: 1 [40960/60000 (68%)]\tLoss:222.458389\n",
      "Train Epoch: 1 [42240/60000 (70%)]\tLoss:224.217773\n",
      "Train Epoch: 1 [43520/60000 (72%)]\tLoss:218.670425\n",
      "Train Epoch: 1 [44800/60000 (75%)]\tLoss:211.459351\n",
      "Train Epoch: 1 [46080/60000 (77%)]\tLoss:214.266739\n",
      "Train Epoch: 1 [47360/60000 (79%)]\tLoss:207.735199\n",
      "Train Epoch: 1 [48640/60000 (81%)]\tLoss:206.659439\n",
      "Train Epoch: 1 [49920/60000 (83%)]\tLoss:205.254608\n",
      "Train Epoch: 1 [51200/60000 (85%)]\tLoss:204.999649\n",
      "Train Epoch: 1 [52480/60000 (87%)]\tLoss:204.978149\n",
      "Train Epoch: 1 [53760/60000 (90%)]\tLoss:209.497757\n",
      "Train Epoch: 1 [55040/60000 (92%)]\tLoss:195.477478\n",
      "Train Epoch: 1 [56320/60000 (94%)]\tLoss:209.405548\n",
      "Train Epoch: 1 [57600/60000 (96%)]\tLoss:201.871033\n",
      "Train Epoch: 1 [58880/60000 (98%)]\tLoss:201.215210\n",
      "[Epoch]:1, Average Loss:260.1906\n",
      "Test Loss:196.609915\n",
      "Train Epoch: 2 [0/60000 (0%)]\tLoss:189.948013\n",
      "Train Epoch: 2 [1280/60000 (2%)]\tLoss:187.925674\n",
      "Train Epoch: 2 [2560/60000 (4%)]\tLoss:192.161652\n",
      "Train Epoch: 2 [3840/60000 (6%)]\tLoss:195.751373\n",
      "Train Epoch: 2 [5120/60000 (9%)]\tLoss:195.773193\n",
      "Train Epoch: 2 [6400/60000 (11%)]\tLoss:199.834732\n",
      "Train Epoch: 2 [7680/60000 (13%)]\tLoss:189.664108\n",
      "Train Epoch: 2 [8960/60000 (15%)]\tLoss:191.761261\n",
      "Train Epoch: 2 [10240/60000 (17%)]\tLoss:191.182526\n",
      "Train Epoch: 2 [11520/60000 (19%)]\tLoss:185.772659\n",
      "Train Epoch: 2 [12800/60000 (21%)]\tLoss:189.040131\n",
      "Train Epoch: 2 [14080/60000 (23%)]\tLoss:180.118622\n",
      "Train Epoch: 2 [15360/60000 (26%)]\tLoss:191.524261\n",
      "Train Epoch: 2 [16640/60000 (28%)]\tLoss:182.181961\n",
      "Train Epoch: 2 [17920/60000 (30%)]\tLoss:186.716202\n",
      "Train Epoch: 2 [19200/60000 (32%)]\tLoss:180.887985\n",
      "Train Epoch: 2 [20480/60000 (34%)]\tLoss:182.848022\n",
      "Train Epoch: 2 [21760/60000 (36%)]\tLoss:180.016037\n",
      "Train Epoch: 2 [23040/60000 (38%)]\tLoss:178.831879\n",
      "Train Epoch: 2 [24320/60000 (41%)]\tLoss:177.614273\n",
      "Train Epoch: 2 [25600/60000 (43%)]\tLoss:175.796341\n",
      "Train Epoch: 2 [26880/60000 (45%)]\tLoss:172.961563\n",
      "Train Epoch: 2 [28160/60000 (47%)]\tLoss:174.548615\n",
      "Train Epoch: 2 [29440/60000 (49%)]\tLoss:176.848328\n",
      "Train Epoch: 2 [30720/60000 (51%)]\tLoss:165.736038\n",
      "Train Epoch: 2 [32000/60000 (53%)]\tLoss:176.726593\n",
      "Train Epoch: 2 [33280/60000 (55%)]\tLoss:169.980881\n",
      "Train Epoch: 2 [34560/60000 (58%)]\tLoss:174.255585\n",
      "Train Epoch: 2 [35840/60000 (60%)]\tLoss:174.244278\n",
      "Train Epoch: 2 [37120/60000 (62%)]\tLoss:175.379196\n",
      "Train Epoch: 2 [38400/60000 (64%)]\tLoss:173.278122\n",
      "Train Epoch: 2 [39680/60000 (66%)]\tLoss:167.965408\n",
      "Train Epoch: 2 [40960/60000 (68%)]\tLoss:172.158020\n",
      "Train Epoch: 2 [42240/60000 (70%)]\tLoss:177.669373\n",
      "Train Epoch: 2 [43520/60000 (72%)]\tLoss:167.128540\n",
      "Train Epoch: 2 [44800/60000 (75%)]\tLoss:171.546631\n",
      "Train Epoch: 2 [46080/60000 (77%)]\tLoss:172.759491\n",
      "Train Epoch: 2 [47360/60000 (79%)]\tLoss:167.569519\n",
      "Train Epoch: 2 [48640/60000 (81%)]\tLoss:169.454834\n",
      "Train Epoch: 2 [49920/60000 (83%)]\tLoss:161.396973\n",
      "Train Epoch: 2 [51200/60000 (85%)]\tLoss:167.571960\n",
      "Train Epoch: 2 [52480/60000 (87%)]\tLoss:164.055740\n",
      "Train Epoch: 2 [53760/60000 (90%)]\tLoss:167.352829\n",
      "Train Epoch: 2 [55040/60000 (92%)]\tLoss:163.232590\n",
      "Train Epoch: 2 [56320/60000 (94%)]\tLoss:163.375244\n",
      "Train Epoch: 2 [57600/60000 (96%)]\tLoss:166.925598\n",
      "Train Epoch: 2 [58880/60000 (98%)]\tLoss:166.316635\n",
      "[Epoch]:2, Average Loss:177.6280\n",
      "Test Loss:163.634380\n",
      "Train Epoch: 3 [0/60000 (0%)]\tLoss:169.310837\n",
      "Train Epoch: 3 [1280/60000 (2%)]\tLoss:163.536133\n",
      "Train Epoch: 3 [2560/60000 (4%)]\tLoss:164.850677\n",
      "Train Epoch: 3 [3840/60000 (6%)]\tLoss:167.197113\n",
      "Train Epoch: 3 [5120/60000 (9%)]\tLoss:164.936371\n",
      "Train Epoch: 3 [6400/60000 (11%)]\tLoss:160.356750\n",
      "Train Epoch: 3 [7680/60000 (13%)]\tLoss:162.114349\n",
      "Train Epoch: 3 [8960/60000 (15%)]\tLoss:165.888397\n",
      "Train Epoch: 3 [10240/60000 (17%)]\tLoss:163.248688\n",
      "Train Epoch: 3 [11520/60000 (19%)]\tLoss:157.241608\n",
      "Train Epoch: 3 [12800/60000 (21%)]\tLoss:160.160706\n",
      "Train Epoch: 3 [14080/60000 (23%)]\tLoss:169.768219\n",
      "Train Epoch: 3 [15360/60000 (26%)]\tLoss:162.919159\n",
      "Train Epoch: 3 [16640/60000 (28%)]\tLoss:158.439789\n",
      "Train Epoch: 3 [17920/60000 (30%)]\tLoss:164.025208\n",
      "Train Epoch: 3 [19200/60000 (32%)]\tLoss:155.336136\n",
      "Train Epoch: 3 [20480/60000 (34%)]\tLoss:160.147888\n",
      "Train Epoch: 3 [21760/60000 (36%)]\tLoss:148.816559\n",
      "Train Epoch: 3 [23040/60000 (38%)]\tLoss:160.481842\n",
      "Train Epoch: 3 [24320/60000 (41%)]\tLoss:158.258011\n",
      "Train Epoch: 3 [25600/60000 (43%)]\tLoss:155.170319\n",
      "Train Epoch: 3 [26880/60000 (45%)]\tLoss:155.929520\n",
      "Train Epoch: 3 [28160/60000 (47%)]\tLoss:153.098083\n",
      "Train Epoch: 3 [29440/60000 (49%)]\tLoss:159.948700\n",
      "Train Epoch: 3 [30720/60000 (51%)]\tLoss:153.067291\n",
      "Train Epoch: 3 [32000/60000 (53%)]\tLoss:150.436829\n",
      "Train Epoch: 3 [33280/60000 (55%)]\tLoss:156.113464\n",
      "Train Epoch: 3 [34560/60000 (58%)]\tLoss:155.698685\n",
      "Train Epoch: 3 [35840/60000 (60%)]\tLoss:158.031342\n",
      "Train Epoch: 3 [37120/60000 (62%)]\tLoss:158.150879\n",
      "Train Epoch: 3 [38400/60000 (64%)]\tLoss:155.013245\n",
      "Train Epoch: 3 [39680/60000 (66%)]\tLoss:149.268784\n",
      "Train Epoch: 3 [40960/60000 (68%)]\tLoss:151.847168\n",
      "Train Epoch: 3 [42240/60000 (70%)]\tLoss:150.281631\n",
      "Train Epoch: 3 [43520/60000 (72%)]\tLoss:158.695099\n",
      "Train Epoch: 3 [44800/60000 (75%)]\tLoss:148.506577\n",
      "Train Epoch: 3 [46080/60000 (77%)]\tLoss:149.271515\n",
      "Train Epoch: 3 [47360/60000 (79%)]\tLoss:148.598358\n",
      "Train Epoch: 3 [48640/60000 (81%)]\tLoss:152.726089\n",
      "Train Epoch: 3 [49920/60000 (83%)]\tLoss:153.721161\n",
      "Train Epoch: 3 [51200/60000 (85%)]\tLoss:155.577271\n",
      "Train Epoch: 3 [52480/60000 (87%)]\tLoss:150.697052\n",
      "Train Epoch: 3 [53760/60000 (90%)]\tLoss:151.010086\n",
      "Train Epoch: 3 [55040/60000 (92%)]\tLoss:147.639053\n",
      "Train Epoch: 3 [56320/60000 (94%)]\tLoss:148.937500\n",
      "Train Epoch: 3 [57600/60000 (96%)]\tLoss:152.037689\n",
      "Train Epoch: 3 [58880/60000 (98%)]\tLoss:149.301651\n",
      "[Epoch]:3, Average Loss:157.1286\n",
      "Test Loss:149.357027\n"
     ]
    }
   ],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    for epoch in range(1, args.epochs +1):\n",
    "        train(epoch)\n",
    "        test(epoch)\n",
    "        # Test Decoder\n",
    "        with torch.no_grad():\n",
    "            sample = torch.randn(64, 20).to(device)\n",
    "            sample = model.decode(sample).cpu()\n",
    "            save_image(sample.view(64, 1, 28, 28),\n",
    "                './results/sample_'+str(epoch)+'.png')       "
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "82ebbbba6002e52ffd5c3db618d60685134ba8ae181eb7df3dc9e0c9fc17132c"
  },
  "kernelspec": {
   "display_name": "Python 3.9.13 ('MAC-GPU')",
   "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.9.13"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}