v2_new.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ae588f3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torchinfo import summary\n",
    "from utils import *\n",
    "from model_training import DoodleDataset, RealDataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9c258fff",
   "metadata": {},
   "outputs": [],
   "source": [
    "class V2ConvNet(nn.Module):\n",
    "    def __init__(self, in_c, \n",
    "                 num_classes, \n",
    "                 channel_list=[64, 128, 192, 256, 512], \n",
    "                 pool_option=(1,1), \n",
    "                 hidden=256, \n",
    "                 dropout=0.2, \n",
    "                 add_layers=False):\n",
    "        super().__init__()\n",
    "        \n",
    "        layer1 = nn.Conv2d(in_c, channel_list[0], kernel_size=3)\n",
    "        layer2 = nn.Conv2d(channel_list[0], channel_list[0], kernel_size=3)\n",
    "        layers = [layer1, layer2]\n",
    "        \n",
    "        for i in range(1, len(channel_list)):\n",
    "            layers.append(\n",
    "                nn.Conv2d(channel_list[i-1], channel_list[i], kernel_size=3, stride=2, padding=1, bias=True)\n",
    "            )\n",
    "            layers.append(\n",
    "                nn.Conv2d(channel_list[i], channel_list[i], kernel_size=3, stride=2, padding=1, bias=True)\n",
    "            )\n",
    "            layers.append(\n",
    "                nn.BatchNorm2d(channel_list[i])\n",
    "            )\n",
    "            layers.append(\n",
    "                nn.Dropout(dropout)\n",
    "            )\n",
    "            layers.append(nn.ReLU())\n",
    "            \n",
    "        self.conv = nn.Sequential(*layers)\n",
    "        \n",
    "        self.flatten = nn.AdaptiveAvgPool2d(pool_option)\n",
    "            \n",
    "        self.fc = nn.Sequential(*[\n",
    "            nn.Linear(pool_option[0] * pool_option[1] * channel_list[-1], hidden),\n",
    "            nn.Linear(hidden, num_classes)\n",
    "        ])\n",
    "\n",
    "    def forward(self, x, return_feats=False):\n",
    "        feats = self.conv(x)\n",
    "        x = x.view(x.size(0), 512, -1).mean(2)\n",
    "        x = self.fc(x)\n",
    "\n",
    "        if return_feats:\n",
    "            return x, feats\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c0082e7c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([100, 9])\n",
      "=================================================================\n",
      "Layer (type:depth-idx)                   Param #\n",
      "=================================================================\n",
      "V2ConvNet                                --\n",
      "├─Sequential: 1-1                        --\n",
      "│    └─Conv2d: 2-1                       1,792\n",
      "│    └─Conv2d: 2-2                       36,928\n",
      "│    └─Conv2d: 2-3                       73,856\n",
      "│    └─Conv2d: 2-4                       147,584\n",
      "│    └─BatchNorm2d: 2-5                  256\n",
      "│    └─Dropout: 2-6                      --\n",
      "│    └─ReLU: 2-7                         --\n",
      "│    └─Conv2d: 2-8                       221,376\n",
      "│    └─Conv2d: 2-9                       331,968\n",
      "│    └─BatchNorm2d: 2-10                 384\n",
      "│    └─Dropout: 2-11                     --\n",
      "│    └─ReLU: 2-12                        --\n",
      "│    └─Conv2d: 2-13                      442,624\n",
      "│    └─Conv2d: 2-14                      590,080\n",
      "│    └─BatchNorm2d: 2-15                 512\n",
      "│    └─Dropout: 2-16                     --\n",
      "│    └─ReLU: 2-17                        --\n",
      "│    └─Conv2d: 2-18                      1,180,160\n",
      "│    └─Conv2d: 2-19                      2,359,808\n",
      "│    └─BatchNorm2d: 2-20                 1,024\n",
      "│    └─Dropout: 2-21                     --\n",
      "│    └─ReLU: 2-22                        --\n",
      "├─AdaptiveAvgPool2d: 1-2                 --\n",
      "├─Sequential: 1-3                        --\n",
      "│    └─Linear: 2-23                      131,328\n",
      "│    └─Linear: 2-24                      2,313\n",
      "=================================================================\n",
      "Total params: 5,521,993\n",
      "Trainable params: 5,521,993\n",
      "Non-trainable params: 0\n",
      "=================================================================\n"
     ]
    }
   ],
   "source": [
    "from torchinfo import summary\n",
    "\n",
    "x = torch.rand(100, 3, 64, 64)\n",
    "net = V2ConvNet(3, 9, [64, 128, 192, 256, 512])\n",
    "y = net(x)\n",
    "print (y.shape)\n",
    "print (summary(net))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2623cf01",
   "metadata": {},
   "outputs": [],
   "source": [
    "real_train_set = RealDataset(train=True)\n",
    "real_val_set = RealDataset(train=False)\n",
    "\n",
    "real_train_loader = torch.utils.data.DataLoader(real_train_set, batch_size=512, shuffle=True)\n",
    "real_val_loader = torch.utils.data.DataLoader(real_val_set, batch_size=512)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1cabcffa",
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = nn.CrossEntropyLoss()\n",
    "model = V2ConvNet(3, 9, [64, 128, 192, 512])\n",
    "optim = torch.optim.AdamW(model.parameters())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "51ea5f38",
   "metadata": {},
   "outputs": [],
   "source": [
    "fix_seed(0)  # zero seed by default\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"2, 3\"\n",
    "\n",
    "model = nn.DataParallel(model).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "b362656c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_accuracy(pred, label):\n",
    "    pred, label = pred.cpu(), label.cpu()\n",
    "    return (pred.argmax(1) == label).sum().item() / len(label)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "a0a1639e",
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "CUDA out of memory. Tried to allocate 20.00 MiB (GPU 3; 23.70 GiB total capacity; 0 bytes already allocated; 20.56 MiB free; 0 bytes reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_2680293/1767585841.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     10\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     11\u001b[0m         \u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcuda\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcuda\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 12\u001b[0;31m         \u001b[0mpred\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     13\u001b[0m         \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcriterion\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpred\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     14\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/envs/Rish4243/lib/python3.7/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1108\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1109\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1110\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1111\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1112\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/envs/Rish4243/lib/python3.7/site-packages/torch/nn/parallel/data_parallel.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, *inputs, **kwargs)\u001b[0m\n\u001b[1;32m    156\u001b[0m                                        \"them on device: {}\".format(self.src_device_obj, t.device))\n\u001b[1;32m    157\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 158\u001b[0;31m             \u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwargs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mscatter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdevice_ids\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    159\u001b[0m             \u001b[0;31m# for forward function without any inputs, empty list and dict will be created\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    160\u001b[0m             \u001b[0;31m# so the module can be executed on one device which is the first one in device_ids\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/envs/Rish4243/lib/python3.7/site-packages/torch/nn/parallel/data_parallel.py\u001b[0m in \u001b[0;36mscatter\u001b[0;34m(self, inputs, kwargs, device_ids)\u001b[0m\n\u001b[1;32m    173\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    174\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mscatter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice_ids\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 175\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mscatter_kwargs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice_ids\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdim\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    176\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    177\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mparallel_apply\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreplicas\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/envs/Rish4243/lib/python3.7/site-packages/torch/nn/parallel/scatter_gather.py\u001b[0m in \u001b[0;36mscatter_kwargs\u001b[0;34m(inputs, kwargs, target_gpus, dim)\u001b[0m\n\u001b[1;32m     42\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mscatter_kwargs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtarget_gpus\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     43\u001b[0m     \u001b[0;34mr\"\"\"Scatter with support for kwargs dictionary\"\"\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 44\u001b[0;31m     \u001b[0minputs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mscatter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtarget_gpus\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdim\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0minputs\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     45\u001b[0m     \u001b[0mkwargs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mscatter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtarget_gpus\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdim\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mkwargs\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     46\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m<\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/envs/Rish4243/lib/python3.7/site-packages/torch/nn/parallel/scatter_gather.py\u001b[0m in \u001b[0;36mscatter\u001b[0;34m(inputs, target_gpus, dim)\u001b[0m\n\u001b[1;32m     34\u001b[0m     \u001b[0;31m# None, clearing the cell\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     35\u001b[0m     \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 36\u001b[0;31m         \u001b[0mres\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mscatter_map\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     37\u001b[0m     \u001b[0;32mfinally\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     38\u001b[0m         \u001b[0mscatter_map\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/envs/Rish4243/lib/python3.7/site-packages/torch/nn/parallel/scatter_gather.py\u001b[0m in \u001b[0;36mscatter_map\u001b[0;34m(obj)\u001b[0m\n\u001b[1;32m     21\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mtype\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0margs\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mzip\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mmap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mscatter_map\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mobj\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     22\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtuple\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m>\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 23\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mlist\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mzip\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mmap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mscatter_map\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mobj\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     24\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlist\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m>\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     25\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mlist\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mzip\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mmap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mscatter_map\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mobj\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/envs/Rish4243/lib/python3.7/site-packages/torch/nn/parallel/scatter_gather.py\u001b[0m in \u001b[0;36mscatter_map\u001b[0;34m(obj)\u001b[0m\n\u001b[1;32m     17\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mscatter_map\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     18\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTensor\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 19\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mScatter\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mapply\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtarget_gpus\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdim\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mobj\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     20\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mis_namedtuple\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     21\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mtype\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0margs\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mzip\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mmap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mscatter_map\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mobj\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/envs/Rish4243/lib/python3.7/site-packages/torch/nn/parallel/_functions.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(ctx, target_gpus, chunk_sizes, dim, input)\u001b[0m\n\u001b[1;32m     94\u001b[0m             \u001b[0;31m# Perform CPU to GPU copies in a background stream\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     95\u001b[0m             \u001b[0mstreams\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0m_get_stream\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mdevice\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mtarget_gpus\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 96\u001b[0;31m         \u001b[0moutputs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcomm\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mscatter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtarget_gpus\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mchunk_sizes\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mctx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdim\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstreams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     97\u001b[0m         \u001b[0;31m# Synchronize with the copy stream\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     98\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mstreams\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/envs/Rish4243/lib/python3.7/site-packages/torch/nn/parallel/comm.py\u001b[0m in \u001b[0;36mscatter\u001b[0;34m(tensor, devices, chunk_sizes, dim, streams, out)\u001b[0m\n\u001b[1;32m    187\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mout\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    188\u001b[0m         \u001b[0mdevices\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0m_get_device_index\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0md\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0md\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mdevices\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 189\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mtuple\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_C\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_scatter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtensor\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevices\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mchunk_sizes\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdim\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstreams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    190\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    191\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mdevices\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: CUDA out of memory. Tried to allocate 20.00 MiB (GPU 3; 23.70 GiB total capacity; 0 bytes already allocated; 20.56 MiB free; 0 bytes reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF"
     ]
    }
   ],
   "source": [
    "epochs= 20\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    total_loss = 0\n",
    "    model.train()\n",
    "    count = 0\n",
    "    for idx, (x, y) in enumerate(real_train_loader):\n",
    "        count += 1\n",
    "        optim.zero_grad()\n",
    "        \n",
    "        x, y = x.cuda(), y.cuda()\n",
    "        pred = model(x)\n",
    "        loss = criterion(pred, y)\n",
    "\n",
    "        total_loss += loss.detach().cpu().item()\n",
    "        loss.backward()\n",
    "        optim.step()\n",
    "    \n",
    "    epoch_loss = total_loss / count\n",
    "    \n",
    "    total_val_acc = 0\n",
    "    with torch.no_grad():\n",
    "        model.eval()\n",
    "        count = 0\n",
    "        for idx, (x, y) in enumerate(real_val_loader):\n",
    "            count += 1\n",
    "            pred = model(x)\n",
    "            val_acc = get_accuracy(pred, y)\n",
    "            total_val_acc += val_acc\n",
    "            \n",
    "        avg_val_acc = total_val_acc / count\n",
    "        \n",
    "    print (\"Epoch: {}, Training Loss: {}, Avg. Validation Accuracy: {}\".format(epoch, epoch_loss, avg_val_acc))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}