NODE_MSTAR_Classifier.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import argparse\n",
    "import logging\n",
    "import time\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import h5py\n",
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format = 'retina'\n",
    "\n",
    "from sklearn.preprocessing import LabelEncoder\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils import data\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "from torchdiffeq import odeint_adjoint as odeint"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# GPU device\n",
    "device = 'cuda:0'\n",
    "\n",
    "# Downsampling method to use\n",
    "# conv: convolutional layers\n",
    "# res: residual layers\n",
    "downsampling_method = 'res'\n",
    "num_channels = 64\n",
    "num_epochs = 100\n",
    "\n",
    "batch_time = 5\n",
    "batch_size = 20\n",
    "total_data = 5000\n",
    "learning_rate = 1e-4\n",
    "\n",
    "# Use CUDA when available\n",
    "if torch.cuda.is_available():\n",
    "    dtype = torch.cuda.FloatTensor\n",
    "    torch.backends.cudnn.benchmark=True\n",
    "else:\n",
    "    dtype = torch.FloatTensor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building Blocks and Definition for an ODENet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [],
   "source": [
    "def conv3x3(in_planes, out_planes, stride=1):\n",
    "    \"\"\"3x3 convolution with padding\"\"\"\n",
    "    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)\n",
    "\n",
    "def conv1x1(in_planes, out_planes, stride=1):\n",
    "    \"\"\"1x1 convolution\"\"\"\n",
    "    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)\n",
    "\n",
    "def norm(dim):\n",
    "    return nn.GroupNorm(min(32, dim), dim)\n",
    "\n",
    "class ConcatConv2d(nn.Module):\n",
    "\n",
    "    def __init__(self, dim_in, dim_out, ksize=3, stride=1, padding=0, dilation=1, groups=1, bias=True, transpose=False):\n",
    "        super(ConcatConv2d, self).__init__()\n",
    "        module = nn.ConvTranspose2d if transpose else nn.Conv2d\n",
    "        self._layer = module(\n",
    "            dim_in + 1, dim_out, kernel_size=ksize, stride=stride, padding=padding, dilation=dilation, groups=groups,\n",
    "            bias=bias\n",
    "        )\n",
    "\n",
    "    def forward(self, t, x):\n",
    "        tt = torch.ones_like(x[:, :1, :, :]) * t\n",
    "        ttx = torch.cat([tt, x], 1)\n",
    "        return self._layer(ttx)\n",
    "    \n",
    "class Flatten(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super(Flatten, self).__init__()\n",
    "\n",
    "    def forward(self, x):\n",
    "        shape = torch.prod(torch.tensor(x.shape[1:])).item()\n",
    "        return x.view(-1, shape)\n",
    "\n",
    "class ResBlock(nn.Module):\n",
    "    expansion = 1\n",
    "\n",
    "    def __init__(self, inplanes, planes, stride=1, downsample=None):\n",
    "        super(ResBlock, self).__init__()\n",
    "        self.norm1 = norm(inplanes)\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        self.downsample = downsample\n",
    "        self.conv1 = conv3x3(inplanes, planes, stride)\n",
    "        self.norm2 = norm(planes)\n",
    "        self.conv2 = conv3x3(planes, planes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        shortcut = x\n",
    "\n",
    "        out = self.relu(self.norm1(x))\n",
    "\n",
    "        if self.downsample is not None:\n",
    "            shortcut = self.downsample(out)\n",
    "\n",
    "        out = self.conv1(out)\n",
    "        out = self.norm2(out)\n",
    "        out = self.relu(out)\n",
    "        out = self.conv2(out)\n",
    "\n",
    "        return out + shortcut\n",
    "\n",
    "class ODEBlock(nn.Module):\n",
    "\n",
    "    def __init__(self, odefunc):\n",
    "        super(ODEBlock, self).__init__()\n",
    "        self.odefunc = odefunc\n",
    "        self.integration_time = torch.tensor([0, 1]).float()\n",
    "\n",
    "    def forward(self, x):\n",
    "        self.integration_time = self.integration_time.type_as(x)\n",
    "        out = odeint(self.odefunc, x, self.integration_time, rtol=args.tol, atol=args.tol)\n",
    "        return out[1]\n",
    "\n",
    "    @property\n",
    "    def nfe(self):\n",
    "        return self.odefunc.nfe\n",
    "\n",
    "    @nfe.setter\n",
    "    def nfe(self, value):\n",
    "        self.odefunc.nfe = value\n",
    "        \n",
    "class ODENet(nn.Module):\n",
    "    \n",
    "    def __init__(self, odefunc):\n",
    "        super(ODENet, self).__init__()\n",
    "        \n",
    "        if downsampling_method is 'conv':\n",
    "            self.downsampling_layers = [\n",
    "                nn.Conv2d(1, num_channels, 3, 1),\n",
    "                norm(num_channels),\n",
    "                nn.ReLU(inplace=True),\n",
    "                nn.Conv2d(num_channels, num_channels, 4, 2, 1),\n",
    "                norm(num_channels),\n",
    "                nn.ReLU(inplace=True),\n",
    "                nn.Conv2d(num_channels, num_channels, 4, 2, 1),\n",
    "            ]\n",
    "        elif downsampling_method is 'res':\n",
    "            self.downsampling_layers = [\n",
    "                nn.Conv2d(1, num_channels, 3, 1),\n",
    "                ResBlock(num_channels, num_channels, stride=2, downsample=conv1x1(num_channels, num_channels, 2)),\n",
    "                ResBlock(num_channels, num_channels, stride=2, downsample=conv1x1(num_channels, num_channels, 2)),\n",
    "            ]\n",
    "\n",
    "        self.feature_layers = [ODEBlock(odefunc(num_channels))]\n",
    "        \n",
    "        self.fc_layers = [norm(64), nn.ReLU(inplace=True), nn.AdaptiveAvgPool2d((1, 1)), \n",
    "                          Flatten(), nn.Linear(64, 10)]\n",
    "        \n",
    "        self._layers = nn.Sequential(\n",
    "                        *self.downsampling_layers,\n",
    "                        *self.feature_layers,\n",
    "                        *self.fc_layers)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self._layers(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define different ODE Function Blocks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ODEClassificationFunc(nn.Module):\n",
    "\n",
    "    def __init__(self, dim):\n",
    "        super(ODEClassificationFunc, self).__init__()\n",
    "        self.norm1 = norm(dim)\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        self.conv1 = ConcatConv2d(dim, dim, 3, 1, 1)\n",
    "        self.norm2 = norm(dim)\n",
    "        self.conv2 = ConcatConv2d(dim, dim, 3, 1, 1)\n",
    "        self.norm3 = norm(dim)\n",
    "        self.nfe = 0\n",
    "\n",
    "    def forward(self, t, x):\n",
    "        self.nfe += 1\n",
    "        out = self.norm1(x)\n",
    "        out = self.relu(out)\n",
    "        out = self.conv1(t, out)\n",
    "        out = self.norm2(out)\n",
    "        out = self.relu(out)\n",
    "        out = self.conv2(t, out)\n",
    "        out = self.norm3(out)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Loading\n",
    "\n",
    "Loading of SAR images and their respective labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "class MSTARDataset(data.Dataset):\n",
    "    'Characterizes a dataset for PyTorch'\n",
    "    def __init__(self, file_name, transform=None):\n",
    "        'Initialization'\n",
    "        super(MSTARDataset, self).__init__()\n",
    "        self.file = h5py.File(file_name, 'r')\n",
    "        self.transform = transform\n",
    "        self.length = self.file['images'].shape[0]\n",
    "\n",
    "    def __len__(self):\n",
    "        'Denotes the total number of samples'\n",
    "        return self.length\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        'Generates one sample of data'\n",
    "        # Select sample\n",
    "        image = self.file['images'][index,:,:].astype('float32')\n",
    "        encoded_label = self.file['encoded_labels'][index]\n",
    "        #label = self.file['labels'][index]\n",
    "    \n",
    "        if self.transform is not None:\n",
    "            self.transform(image)\n",
    "            \n",
    "        return image, encoded_label\n",
    "    \n",
    "def complex_abs(data, dim=-1, keepdim=True):\n",
    "    print(data.shape)\n",
    "    assert data.size(-1) == 2\n",
    "    return (data ** 2).sum(dim=-1, keepdim=True).sqrt()    \n",
    "    \n",
    "transform = transforms.Compose(\n",
    "    [transforms.Lambda(lambda tensor: complex_abs(tensor))])\n",
    "\n",
    "trainloader = data.DataLoader(MSTARDataset('data/mstart_train_data.hdf5', transform), batch_size=batch_size,\n",
    "                                          shuffle=True, num_workers=2)\n",
    "\n",
    "testloader = data.DataLoader(MSTARDataset('data/mstart_test_data.hdf5', transform), batch_size=batch_size,\n",
    "                                         shuffle=False, num_workers=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ODENet(ODEClassificationFunc)\n",
    "model = model.to(device)\n",
    "optimizer = torch.optim.RMSprop(model.parameters(), lr=learning_rate)\n",
    "criterion = nn.CrossEntropyLoss().to(device)\n",
    "end = time.time()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/fallah.5/anaconda3/envs/mstar/lib/python3.6/site-packages/ipykernel_launcher.py:17: ComplexWarning: Casting complex values to real discards the imaginary part\n",
      "/home/fallah.5/anaconda3/envs/mstar/lib/python3.6/site-packages/ipykernel_launcher.py:17: ComplexWarning: Casting complex values to real discards the imaginary part\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[-0.00516798  0.01387577  0.02311549 ... -0.0055943   0.00893019\n",
      "  -0.02152561]\n",
      " [-0.02389601 -0.0281056  -0.00316086 ...  0.04302207 -0.00482417\n",
      "  -0.05653864]\n",
      " [-0.00375925 -0.03050296 -0.00886127 ...  0.00298356 -0.02950739\n",
      "  -0.06167291]\n",
      " ...\n",
      " [ 0.04684278  0.04201766  0.02161201 ... -0.02964322 -0.08212417\n",
      "  -0.10332746]\n",
      " [ 0.03267672  0.0154662  -0.01498783 ... -0.06160478 -0.0968824\n",
      "  -0.13067466]\n",
      " [-0.02965261 -0.023783   -0.05946792 ... -0.05895231 -0.11678518\n",
      "  -0.1156453 ]]\n",
      "[[-0.0087558  -0.03838169 -0.04578689 ... -0.01186337 -0.02341154\n",
      "  -0.01902867]\n",
      " [ 0.02535102  0.02545847  0.01492356 ... -0.01699662 -0.0145135\n",
      "  -0.01648998]\n",
      " [ 0.04728644  0.03865997  0.01525966 ... -0.02407691 -0.01951841\n",
      "   0.01519415]\n",
      " ...\n",
      " [ 0.02741541  0.02692742  0.04524927 ...  0.02085725 -0.0121762\n",
      "  -0.03953116]\n",
      " [ 0.03326316  0.01196783  0.01596735 ...  0.0190334  -0.01358725\n",
      "  -0.04534331]\n",
      " [ 0.02632676 -0.0127333  -0.00112978 ... -0.00319404 -0.01245982\n",
      "  -0.01990991]]\n",
      "[[ 0.0176753  -0.03125764 -0.04882604 ...  0.00512085 -0.02151763\n",
      "  -0.02357347]\n",
      " [ 0.0287043  -0.01270826 -0.02791521 ...  0.00544544 -0.01417927\n",
      "  -0.01524662]\n",
      " [ 0.0088193  -0.03250386 -0.04464448 ... -0.01434198 -0.00221366\n",
      "   0.0154344 ]\n",
      " ...\n",
      " [-0.05204259 -0.02798035 -0.02515512 ... -0.00157751 -0.00802156\n",
      "  -0.04667736]\n",
      " [-0.02658782 -0.02595234 -0.0287043  ...  0.00532489  0.01211589\n",
      "  -0.0230255 ]\n",
      " [-0.00379381 -0.01965713 -0.03483919 ... -0.01750734 -0.00074973\n",
      "  -0.03543623]][[-0.07420393 -0.00160706  0.02140819 ...  0.00051045 -0.0084403\n",
      "  -0.01415535]\n",
      " [-0.01719514  0.00398545  0.02609675 ... -0.00651074 -0.00714727\n",
      "  -0.02723202]\n",
      " [-0.00122389 -0.00396881  0.01724917 ... -0.01111928  0.01286729\n",
      "  -0.01715779]\n",
      " ...\n",
      " [ 0.04292357  0.02764932  0.00081776 ... -0.02248469 -0.03333125\n",
      "  -0.01174983]\n",
      " [-0.01167089 -0.00256081 -0.00331415 ... -0.03769685 -0.02647023\n",
      "   0.01078622]\n",
      " [-0.02008511 -0.02246237  0.00059109 ... -0.03495856 -0.02416844\n",
      "  -0.00069354]]\n",
      "\n"
     ]
    },
    {
     "ename": "TypeError",
     "evalue": "Traceback (most recent call last):\n  File \"/home/fallah.5/anaconda3/envs/mstar/lib/python3.6/site-packages/torch/utils/data/dataloader.py\", line 138, in _worker_loop\n    samples = collate_fn([dataset[i] for i in batch_indices])\n  File \"/home/fallah.5/anaconda3/envs/mstar/lib/python3.6/site-packages/torch/utils/data/dataloader.py\", line 138, in <listcomp>\n    samples = collate_fn([dataset[i] for i in batch_indices])\n  File \"<ipython-input-85-023c1e72d7ee>\", line 22, in __getitem__\n    self.transform(image)\n  File \"/home/fallah.5/anaconda3/envs/mstar/lib/python3.6/site-packages/torchvision/transforms/transforms.py\", line 49, in __call__\n    img = t(img)\n  File \"/home/fallah.5/anaconda3/envs/mstar/lib/python3.6/site-packages/torchvision/transforms/transforms.py\", line 283, in __call__\n    return self.lambd(img)\n  File \"<ipython-input-85-023c1e72d7ee>\", line 32, in <lambda>\n    [transforms.Lambda(lambda tensor: complex_abs(tensor))])\n  File \"<ipython-input-85-023c1e72d7ee>\", line 28, in complex_abs\n    assert data.size(-1) == 2\nTypeError: 'int' object is not callable\n",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-87-490e5a025d4c>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mitr\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnum_epochs\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----> 2\u001b[0;31m     \u001b[0;32mfor\u001b[0m \u001b[0mbatch_count\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch_data\u001b[0m \u001b[0;32min\u001b[0m \u001b[0menumerate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrainloader\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      3\u001b[0m         \u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mzero_grad\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      4\u001b[0m         \u001b[0mimage_batch\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlabel_true\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mbatch_data\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m         \u001b[0mimage_batch\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mimage_batch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/mstar/lib/python3.6/site-packages/torch/utils/data/dataloader.py\u001b[0m in \u001b[0;36m__next__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    635\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreorder_dict\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0midx\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mbatch\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    636\u001b[0m                 \u001b[0;32mcontinue\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 637\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_process_next_batch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbatch\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    638\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    639\u001b[0m     \u001b[0mnext\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m__next__\u001b[0m  \u001b[0;31m# Python 2 compatibility\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/mstar/lib/python3.6/site-packages/torch/utils/data/dataloader.py\u001b[0m in \u001b[0;36m_process_next_batch\u001b[0;34m(self, batch)\u001b[0m\n\u001b[1;32m    656\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_put_indices\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    657\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbatch\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mExceptionWrapper\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--> 658\u001b[0;31m             \u001b[0;32mraise\u001b[0m \u001b[0mbatch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexc_type\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbatch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexc_msg\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    659\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mbatch\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    660\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mTypeError\u001b[0m: Traceback (most recent call last):\n  File \"/home/fallah.5/anaconda3/envs/mstar/lib/python3.6/site-packages/torch/utils/data/dataloader.py\", line 138, in _worker_loop\n    samples = collate_fn([dataset[i] for i in batch_indices])\n  File \"/home/fallah.5/anaconda3/envs/mstar/lib/python3.6/site-packages/torch/utils/data/dataloader.py\", line 138, in <listcomp>\n    samples = collate_fn([dataset[i] for i in batch_indices])\n  File \"<ipython-input-85-023c1e72d7ee>\", line 22, in __getitem__\n    self.transform(image)\n  File \"/home/fallah.5/anaconda3/envs/mstar/lib/python3.6/site-packages/torchvision/transforms/transforms.py\", line 49, in __call__\n    img = t(img)\n  File \"/home/fallah.5/anaconda3/envs/mstar/lib/python3.6/site-packages/torchvision/transforms/transforms.py\", line 283, in __call__\n    return self.lambd(img)\n  File \"<ipython-input-85-023c1e72d7ee>\", line 32, in <lambda>\n    [transforms.Lambda(lambda tensor: complex_abs(tensor))])\n  File \"<ipython-input-85-023c1e72d7ee>\", line 28, in complex_abs\n    assert data.size(-1) == 2\nTypeError: 'int' object is not callable\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 0.04231434  0.07339521  0.05024559 ...  0.01066443 -0.00728734\n",
      "  -0.02187674]\n",
      " [ 0.01291047  0.05053226  0.03470858 ... -0.02431109 -0.04774467\n",
      "  -0.04251335]\n",
      " [-0.03652437 -0.00620858 -0.00162225 ... -0.06948606 -0.07200991\n",
      "  -0.03753733]\n",
      " ...\n",
      " [ 0.01597822 -0.00198883 -0.00359587 ... -0.02968001 -0.0279764\n",
      "  -0.01430927]\n",
      " [-0.02515339  0.00129177  0.00579153 ... -0.05262707 -0.09510272\n",
      "  -0.07532839]\n",
      " [-0.01996377 -0.01881683 -0.01990695 ... -0.03434342 -0.08274035\n",
      "  -0.07062592]]\n"
     ]
    }
   ],
   "source": [
    "for itr in range(0, num_epochs): \n",
    "    model.train()\n",
    "    for batch_count, batch_data in enumerate(trainloader):\n",
    "        optimizer.zero_grad()\n",
    "        image_batch, label_true = batch_data\n",
    "        image_batch = image_batch.to(device)\n",
    "        \n",
    "        label_pred = model(image_batch)\n",
    "        \n",
    "        loss = criterion(label_pred, label_true)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    \n",
    "    total_predict = 0\n",
    "    total_correct = 0\n",
    "    model.eval()\n",
    "    for batch_count, batch_data in enumerate(testloader):\n",
    "        with torch.no_grad():\n",
    "            image_batch, label_true = batch_data\n",
    "            image_batch = image_batch.to(device)\n",
    "            \n",
    "            label_pred = model(image_batch)\n",
    "            \n",
    "            total_predict   += label_true.size(0)\n",
    "            total_correct += (label_pred == label_true).sum().item()\n",
    "\n",
    "    print('Iter {:04d} | Total Correct {:04d} out of {:04d}'.format(itr, total_correct, total_predict))\n",
    "\n",
    "    end = time.time()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (mstar)",
   "language": "python",
   "name": "mstar"
  },
  "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}