notebooks/Chapter-5-code-examples.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import stuff\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "torch.manual_seed(446)\n",
    "np.random.seed(446)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tensors: The analog to numpy arrays  \n",
    "Let's go over the basics of how to use tensors. If you're familiar with numpy, these methods are very similar to those."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x_numpy = [0.1 0.2 0.3] x_torch = tensor([0.1000, 0.2000, 0.3000])\n"
     ]
    }
   ],
   "source": [
    "# we create tensors in a similar way to numpy nd arrays\n",
    "x_numpy = np.array([0.1, 0.2, 0.3])\n",
    "x_torch = torch.tensor([0.1, 0.2, 0.3])\n",
    "print(\"x_numpy =\", x_numpy, \"x_torch =\", x_torch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.1000, 0.2000, 0.3000], dtype=torch.float64) [0.1 0.2 0.3]\n"
     ]
    }
   ],
   "source": [
    "# you can convert between tensors and arrays freely\n",
    "print(torch.from_numpy(x_numpy), x_torch.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x + y\n",
      "[3.1 4.2 5.3] tensor([3.1000, 4.2000, 5.3000])\n"
     ]
    }
   ],
   "source": [
    "# Basic operactions work the same as in numpy\n",
    "y_numpy = np.array([3,4,5.])\n",
    "y_torch = torch.tensor([3,4,5.])\n",
    "print(\"x + y\")\n",
    "print(x_numpy + y_numpy, x_torch + y_torch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "norm\n",
      "0.37416573867739417 tensor(0.3742)\n",
      "\n",
      "mean along the 0th dimension\n",
      "[2. 3.] tensor([2., 3.])\n"
     ]
    }
   ],
   "source": [
    "# many functions that are in numpy are also in pytorch\n",
    "print(\"norm\")\n",
    "print(np.linalg.norm(x_numpy), torch.norm(x_torch))\n",
    "print()\n",
    "\n",
    "# to apply an operation along a dimension,\n",
    "# we use the dim keyword argument instead of axis\n",
    "print(\"mean along the 0th dimension\")\n",
    "x_numpy = np.array([[1,2],[3,4.]])\n",
    "x_torch = torch.tensor([[1,2],[3,4.]])\n",
    "print(np.mean(x_numpy, axis=0), torch.mean(x_torch, dim=0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tensor.view()  \n",
    "This is analogous to numpy.reshape(). Let's use a toy example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([10000, 3, 28, 28])\n",
      "torch.Size([10000, 3, 784])\n",
      "torch.Size([10000, 3, 784])\n"
     ]
    }
   ],
   "source": [
    "N, C, W, H = 10000, 3, 28, 28\n",
    "X = torch.randn((N, C, W, H))\n",
    "\n",
    "print(X.shape)                  # should return [10000, 3, 28, 28]\n",
    "print(X.view(N, C, 784).shape)  # condenses the last 2 dimensions (28 x 28 = 784)\n",
    "print(X.view(-1, C, 784).shape) # passing a -1 automatically choses the correct dimension to reshape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The require_grad keyword  \n",
    "require_grad tells PyTorch to keep track of the [computation graph](https://colah.github.io/posts/2015-08-Backprop/) when using these tensors. In a broad sense computation graphs keep track of the operations done on a set of variables and how the results from those operations rely on their inputs. Setting up these graphs lets you run through backpropagation using the method that we talked about last lecture. So basically, setting require_grad to true will allow you to automatically calculate the gradient through PyTorch!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "c tensor(3., grad_fn=<AddBackward0>)\n",
      "d tensor(2., grad_fn=<AddBackward0>)\n",
      "e tensor(6., grad_fn=<MulBackward0>)\n"
     ]
    }
   ],
   "source": [
    "a = torch.tensor(2.0, requires_grad=True) # we set requires_grad=True to let PyTorch know to keep the graph\n",
    "b = torch.tensor(1.0, requires_grad=True)\n",
    "c = a + b # PyTorch will remember these operations in the computation graph\n",
    "d = b + 1\n",
    "e = c * d\n",
    "print('c', c)\n",
    "print('d', d)\n",
    "print('e', e)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As an example for how we can use this to auto-differentiate functions, lets consider $y = (x-2)^2$.  \n",
    "We could just find the derivative of this function analytically, or we could let PyTorch do it automatically"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Analytical f'(x): tensor([-2.], grad_fn=<MulBackward0>)\n",
      "PyTorch's f'(x): tensor([-2.])\n"
     ]
    }
   ],
   "source": [
    "# define the function\n",
    "def f(x):\n",
    "    return (x-2)**2\n",
    "# define the derivative (we can do so pretty easily)\n",
    "def fp(x):\n",
    "    return 2*(x-2)\n",
    "\n",
    "x = torch.tensor([1.0], requires_grad=True)\n",
    "\n",
    "y = f(x)     # this operation is recorded by PyTorch in the computation graph\n",
    "y.backward() # use the computation graph to calculate gradients at each step\n",
    "\n",
    "print('Analytical f\\'(x):', fp(x))\n",
    "print('PyTorch\\'s f\\'(x):', x.grad)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You get the same answer! So the schematic for **auto differentiation** goes like this:\n",
    "* Define your data as a tensor x  \n",
    "* Do operations on your tensor to get an output o (forward propagation)\n",
    "* Call o.backward()\n",
    "* get the gradient you want with x.grad()  \n",
    "  \n",
    "To further show how you would do this in practice, let's use this auto-differentiation to impliment a basic gradient decent algorithm using the same function as above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "iter,\tx,\tf(x),\tf'(x),\tf'(x) pytorch\n",
      "0,\t5.000,\t9.000,\t6.000,\t6.000\n",
      "1,\t3.500,\t2.250,\t3.000,\t3.000\n",
      "2,\t2.750,\t0.562,\t1.500,\t1.500\n",
      "3,\t2.375,\t0.141,\t0.750,\t0.750\n",
      "4,\t2.188,\t0.035,\t0.375,\t0.375\n",
      "5,\t2.094,\t0.009,\t0.188,\t0.188\n",
      "6,\t2.047,\t0.002,\t0.094,\t0.094\n",
      "7,\t2.023,\t0.001,\t0.047,\t0.047\n",
      "8,\t2.012,\t0.000,\t0.023,\t0.023\n",
      "9,\t2.006,\t0.000,\t0.012,\t0.012\n",
      "10,\t2.003,\t0.000,\t0.006,\t0.006\n",
      "11,\t2.001,\t0.000,\t0.003,\t0.003\n",
      "12,\t2.001,\t0.000,\t0.001,\t0.001\n",
      "13,\t2.000,\t0.000,\t0.001,\t0.001\n",
      "14,\t2.000,\t0.000,\t0.000,\t0.000\n"
     ]
    }
   ],
   "source": [
    "x = torch.tensor([5.0], requires_grad=True) # starting guess\n",
    "step_size = 0.25\n",
    "\n",
    "print('iter,\\tx,\\tf(x),\\tf\\'(x),\\tf\\'(x) pytorch')\n",
    "for i in range(15):\n",
    "    y = f(x)     # this operation gets saved internally\n",
    "    y.backward() # compute the gradient\n",
    "    \n",
    "    print('{},\\t{:.3f},\\t{:.3f},\\t{:.3f},\\t{:.3f}'.format(i, x.item(), f(x).item(), fp(x).item(), x.grad.item()))\n",
    "    x.data = x.data - (step_size * x.grad) # perform a GD update step\n",
    "    \n",
    "    # We need to zero the grad variable since the backward()\n",
    "    # call accumulates the gradients in .grad instead of overwriting.\n",
    "    x.grad.zero_()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Now to Chapter 5 Stuff  \n",
    "This section covers the same stuff as in the slides, but in a more coherent fasion. Much of this code is taken from the book directly. Feel free to play around with things!  \n",
    "## Defining blocks  \n",
    "For more complex networks we want an easier way of defining all the layers than doing so one at a time. To do that, we can use blocks!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.1457, -0.2313,  0.2607,  0.1052, -0.1396, -0.1844,  0.0716, -0.0148,\n",
       "         -0.1052, -0.0598],\n",
       "        [-0.0518, -0.2350,  0.3061, -0.0201, -0.1069, -0.1128,  0.0217,  0.0109,\n",
       "         -0.1015, -0.0378]], grad_fn=<AddmmBackward0>)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# how you would normally make a neural network without blocks\n",
    "net = nn.Sequential(nn.Linear(20, 256), nn.ReLU(), nn.Linear(256, 10))\n",
    "\n",
    "X = torch.rand(2, 20)\n",
    "net(X)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The way you define a block in PyTorch is by defining a class that inherits from nn.Module. Inheriting allows PyTorch to handle a lot of the nitty-gritty details automatically. For example, as long as you aren't defining a new type of operator, PyTorch will autmoatically generate a backpropagation function, so you don't have to! All you have to define for a valid block is a constructor that initializes your block structure, and a forward propagation function. The below cell defines 2 different blocks, and then stiches them together to form a more complex network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(-0.0494, grad_fn=<SumBackward0>)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# first, a standard black with 1 hidden layer and 1 output layer\n",
    "class MLP(nn.Module):\n",
    "    # Declare a layer with model parameters. Here, we declare two fully\n",
    "    # connected layers\n",
    "    def __init__(self):\n",
    "        # Call the constructor of the `MLP` parent class `Module` to perform\n",
    "        # the necessary initialization. In this way, other function arguments\n",
    "        # can also be specified during class instantiation, such as the model\n",
    "        # parameters, `params` (to be described later)\n",
    "        super().__init__()\n",
    "        self.hidden = nn.Linear(20, 256)  # Hidden layer\n",
    "        self.out = nn.Linear(256, 10)  # Output layer\n",
    "\n",
    "    # Define the forward propagation of the model, that is, how to return the\n",
    "    # required model output based on the input `X`\n",
    "    def forward(self, X):\n",
    "        # Note here we use the funtional version of ReLU defined in the\n",
    "        # nn.functional module.\n",
    "        return self.out(F.relu(self.hidden(X)))\n",
    "\n",
    "# Now let's define a more wacky block\n",
    "class FixedHiddenMLP(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        # Random weight parameters that will not compute gradients and\n",
    "        # therefore keep constant during training\n",
    "        self.rand_weight = torch.rand((20, 20), requires_grad=False)\n",
    "        self.linear = nn.Linear(20, 20)\n",
    "\n",
    "    def forward(self, X):\n",
    "        X = self.linear(X)\n",
    "        # Use the created constant parameters, as well as the `relu` and `mm`\n",
    "        # functions\n",
    "        X = F.relu(torch.mm(X, self.rand_weight) + 1)\n",
    "        # Reuse the fully-connected layer. This is equivalent to sharing\n",
    "        # parameters with two fully-connected layers\n",
    "        X = self.linear(X)\n",
    "        # Control flow (this code is arbitrary and just to show that you can do stuff like this\n",
    "        while X.abs().sum() > 1:\n",
    "            X /= 2\n",
    "        return X.sum()\n",
    "\n",
    "# Now let's stitch these 2 blocks together - note that you send blocks as parameters passed to Sequential instead of just layers\n",
    "# So this network has an MLP block, followed by a linear layer, and finally a FixedHiddenMLP block\n",
    "# 20 -> 256 -> 10 -> 20 -> 20\n",
    "net = nn.Sequential(MLP(), nn.Linear(10, 20), FixedHiddenMLP())\n",
    "net(X)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Custom Layers  \n",
    "We can use this method of making blocks to make custom individual layers as well"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-2., -1.,  0.,  1.,  2.])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Let's define a layer that doesn't have any tunable parameters, and just subtracts the mean of the input\n",
    "class CenteredLayer(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "    def forward(self, X):\n",
    "        return X - X.mean()\n",
    "    \n",
    "layer = CenteredLayer()\n",
    "layer(torch.FloatTensor([1, 2, 3, 4, 5]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Parameter containing:\n",
      "tensor([[ 1.2583, -1.0947, -0.2238],\n",
      "        [ 1.7347, -0.1966,  0.9766],\n",
      "        [ 0.0026,  0.7633, -1.1057],\n",
      "        [-0.4654, -1.4138,  1.2072],\n",
      "        [-1.0673, -0.7208,  0.1859]], requires_grad=True)\n",
      "\n",
      "tensor([[0.],\n",
      "        [0.]])\n"
     ]
    }
   ],
   "source": [
    "# Let's try making a layer with parameters this time. This is basically just a normal layer but without an activation function\n",
    "class MyLinear(nn.Module):\n",
    "    def __init__(self, in_units, out_units):\n",
    "        super().__init__()\n",
    "        self.weight = nn.Parameter(torch.randn(in_units, out_units))\n",
    "        self.bias = nn.Parameter(torch.randn(out_units,))\n",
    "    def forward(self, X):\n",
    "        linear = torch.matmul(X, self.weight.data) + self.bias.data\n",
    "        return F.relu(linear)\n",
    "\n",
    "# Notice you have to pass in the size of the inputs and outputs\n",
    "linear = MyLinear(5, 3)\n",
    "print(linear.weight)\n",
    "print()\n",
    "\n",
    "# Just like beofre with blocks, you can combine these custom layers together to make larger networks\n",
    "net = nn.Sequential(MyLinear(64, 8), MyLinear(8, 1))\n",
    "print(net(torch.rand(2, 64)))\n",
    "\n",
    "# EXERSIZE: Try making your own custom layer!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Accessing Parameters  \n",
    "Sometimes you want to be able to access your parameters, like for debugging or for saving your network. There are a few ways you can do that"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sequential(\n",
      "  (0): Linear(in_features=4, out_features=8, bias=True)\n",
      "  (1): ReLU()\n",
      "  (2): Linear(in_features=8, out_features=4, bias=True)\n",
      "  (3): ReLU()\n",
      "  (4): Linear(in_features=4, out_features=2, bias=True)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "# Let's start by defining a network we can manipulate the parameters of\n",
    "params_net = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 4), nn.ReLU(), nn.Linear(4, 2))\n",
    "\n",
    "# to see what this network looks like, we can just print it\n",
    "print(params_net)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OrderedDict([('weight', tensor([[-0.0866, -0.0627, -0.0664, -0.1503],\n",
      "        [-0.2659, -0.0934,  0.1608,  0.3442]])), ('bias', tensor([0.0628, 0.0727]))])\n",
      "tensor([0.0628, 0.0727])\n",
      "\n",
      "<class 'torch.nn.parameter.Parameter'>\n",
      "Parameter containing:\n",
      "tensor([0.0628, 0.0727], requires_grad=True)\n",
      "tensor([0.0628, 0.0727])\n",
      "\n",
      "tensor([0.0628, 0.0727])\n"
     ]
    }
   ],
   "source": [
    "# Let's say we want the bias from the last layer\n",
    "# We can query the layer you want using state_dict() then specify the bias\n",
    "print(params_net[4].state_dict())\n",
    "print(params_net[4].state_dict()[\"bias\"])\n",
    "print()\n",
    "\n",
    "# We could also just access the bias directly\n",
    "print(type(params_net[4].bias))\n",
    "print(params_net[4].bias)\n",
    "print(params_net[4].bias.data) # using .data will get you just the values, no overhead\n",
    "print()\n",
    "\n",
    "# Finally, you can query the entire net, then specify you want the bias from the 4th entry\n",
    "print(params_net.state_dict()['4.bias'].data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initializing Parameters  \n",
    "We've learned that the way you initialize your network is very inportant, but how do you actually do that? You may have noticed I haven't bothered specifying how parameters are initialized in the above examples, and that's because PyTorch automatically does so for us using a normal distribution with $\\sigma = 1 / \\sqrt{N_{weights}}$. If we want more control over initialization, we can do it ourselves by using some common in-built initilization functions, or by creating our own function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([-0.0113,  0.0058, -0.0023,  0.0110]), tensor(0.))"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Standard normal distribution\n",
    "def init_normal(m):\n",
    "    if type(m) == nn.Linear:\n",
    "        nn.init.normal_(m.weight, mean=0, std=0.01)\n",
    "        nn.init.zeros_(m.bias)\n",
    "\n",
    "# let's test this using the network from above\n",
    "params_net.apply(init_normal)\n",
    "params_net[0].weight.data[0], params_net[0].bias.data[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([-0.0128, -0.0364, -0.0721,  0.4716])\n",
      "tensor([[1., 1., 1., 1., 1., 1., 1., 1.],\n",
      "        [1., 1., 1., 1., 1., 1., 1., 1.],\n",
      "        [1., 1., 1., 1., 1., 1., 1., 1.],\n",
      "        [1., 1., 1., 1., 1., 1., 1., 1.]])\n"
     ]
    }
   ],
   "source": [
    "# You can use an xavier distribution too\n",
    "def xavier(m):\n",
    "    if type(m) == nn.Linear:\n",
    "        nn.init.xavier_uniform_(m.weight)\n",
    "# Or, you can even just make all the weights constant\n",
    "def init_constant(m):\n",
    "    if type(m) == nn.Linear:\n",
    "        nn.init.constant_(m.weight, 1)\n",
    "        nn.init.zeros_(m.bias)\n",
    "\n",
    "params_net[0].apply(xavier)\n",
    "params_net[2].apply(init_constant)\n",
    "print(params_net[0].weight.data[0])\n",
    "print(params_net[2].weight.data)\n",
    "\n",
    "# EXERSIZE: Try making your own parameter initilization function! "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sharing Parameters  \n",
    "Sometimes we may want multiple layers to share the same weights. We can make that happen pretty easily"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([True, True, True, True, True, True, True, True])\n",
      "tensor([True, True, True, True, True, True, True, True])\n"
     ]
    }
   ],
   "source": [
    "# We need to give the shared layer a name so that we can refer to its\n",
    "# parameters\n",
    "shared = nn.Linear(8, 8)\n",
    "net = nn.Sequential(nn.Linear(4, 8), nn.ReLU(),\n",
    "                    shared, nn.ReLU(),\n",
    "                    shared, nn.ReLU(),\n",
    "                    nn.Linear(8, 1))\n",
    "net(torch.rand(size=(2, 4)))\n",
    "\n",
    "# Check whether the parameters are the same\n",
    "print(net[2].weight.data[0] == net[4].weight.data[0])\n",
    "net[2].weight.data[0, 0] = 100\n",
    "# Make sure that they are actually the same object rather than just having the\n",
    "# same value\n",
    "print(net[2].weight.data[0] == net[4].weight.data[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Saving and Loading  \n",
    "PyTorch lets you save both individual tensors, and whole networks, to file for checkpointing / saving. The major caveat for saving networks is that it only saves the parameters, and not the structure. So, you'll need to rebuild or copy the structure over in code before loading in the parameters again"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0, 1, 2, 3])\n",
      "tensor([0, 1, 2, 3])\n"
     ]
    }
   ],
   "source": [
    "# Saving tensors is easy - just like numpy!\n",
    "x = torch.arange(4)\n",
    "print(x)\n",
    "torch.save(x, 'x-file')\n",
    "\n",
    "y = torch.load('x-file')\n",
    "print(y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OrderedDict([('weight', tensor([[ 0.0113, -0.0101,  0.0112,  0.0005],\n",
      "        [-0.0008, -0.0005, -0.0174, -0.0270]])), ('bias', tensor([0., 0.]))])\n",
      "OrderedDict([('weight', tensor([[ 0.0113, -0.0101,  0.0112,  0.0005],\n",
      "        [-0.0008, -0.0005, -0.0174, -0.0270]])), ('bias', tensor([0., 0.]))])\n"
     ]
    }
   ],
   "source": [
    "# Let's try saving the network we used above\n",
    "torch.save(params_net.state_dict(), 'p_net.params')\n",
    "\n",
    "# To test loading in these parameters, let's make a clone and load in p_net.params\n",
    "clone_net = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 4), nn.ReLU(), nn.Linear(4, 2))\n",
    "clone_net.load_state_dict(torch.load('p_net.params'))\n",
    "clone_net.eval()\n",
    "\n",
    "# If we've loaded correctly, the print statements below should be the same. Feel free to verify this for different layers!\n",
    "print(params_net[4].state_dict())\n",
    "print(clone_net[4].state_dict())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running on a GPU  \n",
    "This last section requires an NVIDIA GPU and CUDA to run properly. Even if your machine doesn't have a GPU, it might still be useful to take a look at the code "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cpu')"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.device('cpu'), torch.device('cuda')\n",
    "# Here's a helper function that specifies the device as gpu if one exists. If there isn't a GPU, it'll return the cpu device\n",
    "def try_gpu(i=0):  #@save\n",
    "    \"\"\"Return gpu(i) if exists, otherwise return cpu().\"\"\"\n",
    "    if torch.cuda.device_count() >= i + 1:\n",
    "        return torch.device(f'cuda:{i}')\n",
    "    return torch.device('cpu')\n",
    "try_gpu() # If setup correctly this should return device(type=\"cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[1., 1., 1.],\n",
      "        [1., 1., 1.]])\n"
     ]
    }
   ],
   "source": [
    "# Let's try doing something on the GPU (this code will still run if not setup correctly, but on a cpu like normal)\n",
    "X = torch.ones(2, 3, device=try_gpu())\n",
    "Y = torch.rand(2, 3, device=try_gpu())\n",
    "print(X)  # If done right this should specify device=\"cuda\"\n",
    "Z = X + Y # This operation will happen on the GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 0.0261, -0.0924], grad_fn=<AddBackward0>)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Finally, you can transfer networks to GPU if you want\n",
    "GPU_net = params_net.to(device=try_gpu())\n",
    "GPU_net(torch.rand(4)) # This operation will happen on the GPU"
   ]
  },
  {
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}