{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "geographic-argentina",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loading training and test data for 28 x 28 images \n",
    "import torch\n",
    "\n",
    "\n",
    "# Uncomment this section to load data for 112 x 112 images instead\n",
    "data_tr = torch.load(\"training-test-data/training_data_augmented.txt\")\n",
    "target_tr = torch.load(\"training-test-data/training_targets_augmented.txt\")\n",
    "data_test = torch.load(\"training-test-data/test_data_augmented.txt\")\n",
    "target_test = torch.load(\"training-test-data/test_targets_augmented.txt\")\n",
    "\n",
    "num_samples = 128\n",
    "num_samples_test = 100 # number of test samples\n",
    "new_dim = 112 \n",
    "old_dim = 28 # MNIST original dimension\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "every-immigration",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import math\n",
    "from matplotlib.colors import LogNorm\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Two-layer linear convolutional neural network\n",
    "output_channels = 1\n",
    "class Net(nn.Module):\n",
    "    def __init__(self, ker_size1, ker_size2, output_channels):\n",
    "        super(Net, self).__init__()\n",
    "        self.ker_size1 = ker_size1\n",
    "        self.ker_size2 = ker_size2\n",
    "        self.output_channels = output_channels\n",
    "        self.conv1 = nn.Conv2d(1, output_channels, kernel_size=(self.ker_size1, self.ker_size2), bias=False) \n",
    "        self.fc1 = nn.Linear(int(new_dim * new_dim * output_channels), 1, bias=False)\n",
    "\n",
    "\n",
    "    def forward(self, x):\n",
    "        y1 = F.pad(x, (0,self.ker_size2-1,0,self.ker_size1-1), mode='circular') # Circular padding \n",
    "        y1 = self.conv1(y1)\n",
    "        y1 = y1.reshape(y1.size(0), -1)\n",
    "        y1 = self.fc1(y1) \n",
    "        return y1\n",
    "\n",
    "    def initialize(self, initialization_scale):\n",
    "        self.fc1.weight.data.mul_(initialization_scale)\n",
    "        self.conv1.weight.data.mul_(initialization_scale)\n",
    "\n",
    "output = torch.zeros((num_samples, 1))\n",
    "output = output.float()\n",
    "output_test = torch.zeros((num_samples_test, 1))\n",
    "output_test = output.float()\n",
    "\n",
    "\n",
    "# Batch gradient descent\n",
    "def train_minibatch(network, optimizer):\n",
    "    minibatch_size = 32\n",
    "    num_batch = int(num_samples/minibatch_size)\n",
    "    for i in range(num_batch):\n",
    "        network.train()\n",
    "        optimizer.zero_grad()\n",
    "        start_index = i * minibatch_size\n",
    "        end_index = start_index + minibatch_size\n",
    "        output = network(data_tr[start_index:end_index])\n",
    "        loss = torch.sum(torch.exp(-1 * torch.mul(output.flatten(), target_tr[start_index:end_index]))) / minibatch_size\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "# Evaluate training data loss\n",
    "def train_eval(network):\n",
    "    network.eval()\n",
    "    train_loss = 0\n",
    "    correct = 0\n",
    "    with torch.no_grad():\n",
    "        output = network(data_tr)\n",
    "        train_loss = torch.sum(torch.exp(-1 * torch.mul(output.flatten(), target_tr)))\n",
    "        pred = output.apply_(lambda x: 1 if x > 0 else -1)\n",
    "        correct += pred.eq(target_tr.data.view_as(pred)).sum()\n",
    "    train_loss /= num_samples\n",
    "    print('\\nTrain set: Avg. loss: {:.9f}, Accuracy: {}/{} ({:.0f}%)\\n'.format(\n",
    "    train_loss, correct, num_samples,\n",
    "    100. * correct / num_samples))\n",
    "    return train_loss\n",
    "\n",
    "def test(network):\n",
    "    network.eval()\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    with torch.no_grad():\n",
    "        output_test = network(data_test)\n",
    "        test_loss = torch.sum(torch.exp(-1 * torch.mul(output_test.flatten(), target_test)))\n",
    "        pred = output_test.apply_(lambda x: 1 if x > 0 else -1)\n",
    "        correct += pred.eq(target_test.data.view_as(pred)).sum()\n",
    "    test_loss /= num_samples_test\n",
    "    accuracy = 100. * correct / num_samples_test\n",
    "    losses = test_loss\n",
    "    return (accuracy, losses)\n",
    "\n",
    "\n",
    "# Get the information about beta\n",
    "def extract_info(network, show_photo): \n",
    "\n",
    "  # Compute beta for linear CNNs\n",
    "    beta_test = np.zeros((new_dim,new_dim))\n",
    "    for i in range(new_dim):\n",
    "        for j in range(new_dim):\n",
    "            tempimg = torch.zeros((1,1,new_dim, new_dim))\n",
    "            tempimg[0,0,i,j]=1\n",
    "            beta_test[i,j] = network(tempimg)  \n",
    "\n",
    "  # Compute margin\n",
    "    with torch.no_grad():\n",
    "        network.eval()\n",
    "        output_np = np.ndarray.flatten(network(data_tr).data.numpy())\n",
    "        target_np = np.ndarray.flatten(target_tr.data.numpy())\n",
    "        margins = [target_np[i] * output_np[i] for i in range(num_samples)]\n",
    "        min_margin = min(margins) # get the minimum margin for any datapoint \n",
    "\n",
    "\n",
    "    # Compute R(beta)\n",
    "    w1 = network.conv1.weight.detach().numpy()\n",
    "    w2 = network.fc1.weight.detach().numpy()\n",
    "    w1_norm_sq = np.sum(np.square(w1))\n",
    "    w2_norm_sq = np.sum(np.square(w2))\n",
    "    print(w1_norm_sq, w2_norm_sq)\n",
    "    Rbeta = (np.sum(np.square(w1)) + np.sum(np.square(w2))) * np.sqrt(new_dim * new_dim)\n",
    "    Rbeta2 = 2 * new_dim * np.linalg.norm(w1) * np.linalg.norm(w2)\n",
    "\n",
    "    # Normalize by margin \n",
    "    beta_test = beta_test / min_margin # normalize to have margin 1\n",
    "    hat_beta = np.absolute(np.fft.fft2(beta_test,norm='ortho'))\n",
    "    Rbeta = Rbeta / min_margin\n",
    "    Rbeta2 = Rbeta2 / min_margin\n",
    "    \n",
    "    print(\"l2 norm: \" + str(2 * np.sqrt(new_dim * new_dim)* np.linalg.norm(beta_test, ord=\"fro\")))\n",
    "    print(\"l1 norm: \" + str(2 * np.sum(hat_beta)))\n",
    "    print(\"Rbeta: \" + str(Rbeta))\n",
    "    print(\"Rbeta2: \" + str(Rbeta2))\n",
    "\n",
    "    if show_photo:\n",
    "        print(\"Time domain:\")\n",
    "        plt.imshow(np.absolute(beta_test), cmap='gray')\n",
    "        plt.show()\n",
    "        print(\"Frequency domain:\")\n",
    "        plt.imshow(np.absolute(hat_beta), cmap='gray', norm=LogNorm(vmin=0.0001, vmax=0.08))\n",
    "        plt.show()\n",
    "  \n",
    "    return (Rbeta, beta_test)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "external-possibility",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train and extract info about beta\n",
    "import seaborn as sns\n",
    "n_epochs = 100000\n",
    "learning_rate_start = 0.001\n",
    "momentum = 0.3\n",
    "initialization_scale = 0.001\n",
    "\n",
    "from tqdm import tqdm_notebook as tqdm\n",
    "def experiment(ker_size1, ker_size2, output_channels):\n",
    "  # print(class1, class2)\n",
    "    network = Net(ker_size1, ker_size2, output_channels)\n",
    "    network.initialize(initialization_scale)\n",
    "    optimizer =  optim.SGD(network.parameters(), lr=learning_rate_start, momentum=momentum)\n",
    "    print(\"Before training:\")\n",
    "    train_eval(network)\n",
    "    extract_info(network, False)\n",
    "    # test()\n",
    "    lossarray = []\n",
    "    rbetavals = []\n",
    "    ell1s = []\n",
    "    print(\"Start training:\")\n",
    "    for epoch in tqdm(range(1, n_epochs + 1)):\n",
    "        train_minibatch(network, optimizer)\n",
    "        if epoch % 100 == 0:\n",
    "            lossv = train_eval(network)\n",
    "            loss = np.ndarray.flatten(lossv.detach().numpy())[0]\n",
    "            if loss <= 0.000001: # stop at 10^-6 loss \n",
    "                break\n",
    "                \n",
    "#         # After enough epochs, change the learning rate to be higher to expedite convergence\n",
    "        if epoch == 200 == 0:\n",
    "            optimizer =  optim.SGD(network.parameters(), lr=0.005, momentum=momentum)\n",
    "\n",
    "        if epoch == 500:\n",
    "            optimizer = optim.SGD(network.parameters(), lr=0.01, momentum=momentum)\n",
    "\n",
    "        if epoch == 1000:\n",
    "            optimizer = optim.SGD(network.parameters(), lr=0.05, momentum=momentum)\n",
    "\n",
    "        if epoch == 1200:\n",
    "            optimizer = optim.SGD(network.parameters(), lr=0.1, momentum=momentum)\n",
    "\n",
    "        if epoch == 1500:\n",
    "            optimizer = optim.SGD(network.parameters(), lr=0.5, momentum=momentum)\n",
    "\n",
    "        if epoch == 2000:\n",
    "            optimizer = optim.SGD(network.parameters(), lr=1, momentum=momentum)\n",
    "\n",
    "        if epoch == 3000:\n",
    "            optimizer = optim.SGD(network.parameters(), lr=2, momentum=momentum)\n",
    "\n",
    "        if epoch == 4000:\n",
    "            optimizer = optim.SGD(network.parameters(), lr=4, momentum=momentum)\n",
    "\n",
    "        if epoch == 4500:\n",
    "            optimizer = optim.SGD(network.parameters(), lr=10, momentum=momentum)\n",
    "\n",
    "        if epoch == 5000:\n",
    "            optimizer = optim.SGD(network.parameters(), lr=20, momentum=momentum)\n",
    "\n",
    "\n",
    "    print(\"After training:\")\n",
    "    train_eval(network)\n",
    "    (accuracy, losses) = test(network)\n",
    "    print(accuracy, losses)\n",
    "\n",
    "    (rk, beta) = extract_info(network, True)\n",
    "\n",
    "    return (rk, beta)\n",
    "\n",
    "\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "exposed-flexibility",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_experiment(k, T, Cout):\n",
    "    pairs = []\n",
    "    for c in Cout:\n",
    "        pairs.append((k, c))\n",
    "\n",
    "    betas = []\n",
    "    rbetas = []\n",
    "    losses_all = []\n",
    "    rbetavals_all = []\n",
    "    ell1s_all = []\n",
    "\n",
    "    for (k, output_channels) in pairs:\n",
    "        for t in range(T):\n",
    "            print(k, output_channels, t)\n",
    "            (Rbeta, beta) = experiment(k, k, output_channels)\n",
    "            rbetas.append(Rbeta)\n",
    "            betas.append(beta)\n",
    "\n",
    "    # Write data to a CSV\n",
    "    import pandas as pd\n",
    "\n",
    "\n",
    "    # Write betas, losses, and ell1s\n",
    "    for i in range(len(pairs)):\n",
    "        # Write rbetas \n",
    "        rbetas_to_write = []\n",
    "        for t in range(T):\n",
    "            index = i *T + t\n",
    "            rbetas_to_write.append(rbetas[index])\n",
    "            beta = betas[index]\n",
    "            print(str(pairs[i]))\n",
    "            name = \"experiments-data/\" + \"betas-linear-augmented\" + str(pairs[i]) + str(t) +  \".csv\"\n",
    "            pd.DataFrame(beta).to_csv(name, header=False, index=False)\n",
    "        name =  \"experiments-data/\" + str(pairs[i]) + \"rbetas-linear-augmented\" + \".csv\"\n",
    "        pd.DataFrame(rbetas_to_write).to_csv(name, header=False, index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "shaped-elder",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 1 0\n",
      "random initialization linear\n",
      "Before training:\n",
      "\n",
      "Train set: Avg. loss: 0.999999881, Accuracy: 81/128 (63%)\n",
      "\n",
      "1.8110042e-07 3.392222e-07\n",
      "l2 norm: 21.207373156946847\n",
      "l1 norm: 18.650583882871665\n",
      "Rbeta: -22.260130225092954\n",
      "Rbeta2: -21.207373939386088\n",
      "Start training:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<ipython-input-6-93764923a7b6>:22: TqdmDeprecationWarning: This function will be removed in tqdm==5.0.0\n",
      "Please use `tqdm.notebook.tqdm` instead of `tqdm.tqdm_notebook`\n",
      "  for epoch in tqdm(range(1, n_epochs + 1)):\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b9a7431e75344d4a93a0ce81528928b4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/100000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Train set: Avg. loss: 0.877412200, Accuracy: 128/128 (100%)\n",
      "\n",
      "\n",
      "Train set: Avg. loss: 0.078431100, Accuracy: 128/128 (100%)\n",
      "\n",
      "\n",
      "Train set: Avg. loss: 0.031933725, Accuracy: 128/128 (100%)\n",
      "\n",
      "\n",
      "Train set: Avg. loss: 0.018939095, Accuracy: 128/128 (100%)\n",
      "\n",
      "\n",
      "Train set: Avg. loss: 0.013083297, Accuracy: 128/128 (100%)\n",
      "\n",
      "\n",
      "Train set: Avg. loss: 0.002656720, Accuracy: 128/128 (100%)\n",
      "\n",
      "\n",
      "Train set: Avg. loss: 0.001360278, Accuracy: 128/128 (100%)\n",
      "\n",
      "\n",
      "Train set: Avg. loss: 0.000887921, Accuracy: 128/128 (100%)\n",
      "\n",
      "\n",
      "Train set: Avg. loss: 0.000649413, Accuracy: 128/128 (100%)\n",
      "\n",
      "\n",
      "Train set: Avg. loss: 0.000507408, Accuracy: 128/128 (100%)\n",
      "\n",
      "\n",
      "Train set: Avg. loss: 0.000232473, Accuracy: 128/128 (100%)\n",
      "\n",
      "\n",
      "Train set: Avg. loss: 0.000146580, Accuracy: 128/128 (100%)\n",
      "\n",
      "\n",
      "Train set: Avg. loss: 0.000081958, Accuracy: 128/128 (100%)\n",
      "\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-8-24e824b69e16>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mk\u001b[0m \u001b[0;32min\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m3\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m8\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m16\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m28\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----> 6\u001b[0;31m     \u001b[0mrun_experiment\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mk\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mT\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mCout\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      7\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-7-d17e7c5370f5>\u001b[0m in \u001b[0;36mrun_experiment\u001b[0;34m(k, T, Cout)\u001b[0m\n\u001b[1;32m     13\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mt\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mT\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     14\u001b[0m             \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mk\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moutput_channels\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 15\u001b[0;31m             \u001b[0;34m(\u001b[0m\u001b[0mRbeta\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbeta\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mexperiment\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mk\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mk\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moutput_channels\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     16\u001b[0m             \u001b[0mrbetas\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mRbeta\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     17\u001b[0m             \u001b[0mbetas\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbeta\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-6-93764923a7b6>\u001b[0m in \u001b[0;36mexperiment\u001b[0;34m(ker_size1, ker_size2, output_channels)\u001b[0m\n\u001b[1;32m     21\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"Start training:\"\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;32mfor\u001b[0m \u001b[0mepoch\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mtqdm\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mn_epochs\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;36m1\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---> 23\u001b[0;31m         \u001b[0mtrain_minibatch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnetwork\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptimizer\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[0mepoch\u001b[0m \u001b[0;34m%\u001b[0m \u001b[0;36m100\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[0mlossv\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtrain_eval\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnetwork\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-2-c92de4fa0cbe>\u001b[0m in \u001b[0;36mtrain_minibatch\u001b[0;34m(network, optimizer)\u001b[0m\n\u001b[1;32m     51\u001b[0m         \u001b[0moutput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnetwork\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdata_tr\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mstart_index\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0mend_index\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     52\u001b[0m         \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msum\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexp\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmul\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutput\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mflatten\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtarget_tr\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mstart_index\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0mend_index\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[0mminibatch_size\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 53\u001b[0;31m         \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\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     54\u001b[0m         \u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\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     55\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.9/site-packages/torch/tensor.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(self, gradient, retain_graph, create_graph, inputs)\u001b[0m\n\u001b[1;32m    243\u001b[0m                 \u001b[0mcreate_graph\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    244\u001b[0m                 inputs=inputs)\n\u001b[0;32m--> 245\u001b[0;31m         \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mautograd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgradient\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\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    246\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    247\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mregister_hook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhook\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/usr/local/lib/python3.9/site-packages/torch/autograd/__init__.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001b[0m\n\u001b[1;32m    143\u001b[0m         \u001b[0mretain_graph\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    144\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 145\u001b[0;31m     Variable._execution_engine.run_backward(\n\u001b[0m\u001b[1;32m    146\u001b[0m         \u001b[0mtensors\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgrad_tensors_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\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[1;32m    147\u001b[0m         allow_unreachable=True, accumulate_grad=True)  # allow_unreachable flag\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "# Run experiments \n",
    "T = 1\n",
    "Cout = 1\n",
    "\n",
    "for k in [1, 3, 8, 16, 28]:\n",
    "    run_experiment(k, T, Cout)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "looking-typing",
   "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.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
