protein_fcn_prelim2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "get some proteins sequences from one of jeppe's data files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import imp\n",
    "import re\n",
    "from torch.nn.utils.rnn import pad_sequence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_file = '/Users/Deathvoodoo/Documents/openprotein/data/raw/protein_net_testfile.txt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "aa_id_dict = {'A': 0, 'C': 1, 'D': 2, 'E': 3, 'F': 4, 'G': 5, 'H': 6, 'I': 7,\n",
    "              'K': 8, 'L': 9, 'M': 10, 'N': 11, 'P': 12, 'Q': 13, 'R': 14, \n",
    "              'S': 15, 'T': 16, 'V': 17, 'W': 18, 'Y': 19}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def aa_to_onehot(aa_str, aa_to_nr, mask=None):\n",
    "    if mask!=None:\n",
    "        mask_ind = np.asarray([x=='+' for x in mask])*1\n",
    "        mask_ind = np.nonzero(mask_ind)\n",
    "        aa_str = \"\".join([aa_str[x] for x in mask_ind[0]]) # because it gets put in another list\n",
    "    init_array = np.zeros( (len(aa_to_nr.keys()), len(aa_str)) )\n",
    "    for i,j in enumerate(aa_str):\n",
    "        init_array[aa_to_nr[j], i] = 1\n",
    "    return(init_array)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "id_list = []\n",
    "seqs = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(sample_file) as input:\n",
    "    lines = input.readlines()\n",
    "    curr_id = None\n",
    "    for i, line in enumerate(lines):\n",
    "        line = line.strip()\n",
    "        if line == '[ID]':\n",
    "            curr_id = lines[i+1].strip()\n",
    "            id_list.append(curr_id)\n",
    "            seqs[curr_id] = {}\n",
    "            seqs[curr_id]['primary'] = lines[i+3].strip()\n",
    "        if line == '[TERTIARY]':\n",
    "            coords = []\n",
    "            for j in range(3):\n",
    "                coords.append(np.fromstring(lines[i+j+1], sep='\\t'))\n",
    "            seqs[curr_id]['tertiary'] = np.array(coords)\n",
    "        if line == '[MASK]':\n",
    "            seqs[curr_id]['mask'] = lines[i+1].strip()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def filter_coords(coords, mask):\n",
    "    mask = np.array([x=='+' for x in mask])\n",
    "    mask_stretched = np.repeat(mask, 3)\n",
    "    coords_filt = coords[:, mask_stretched]\n",
    "    return(coords_filt)\n",
    "\n",
    "def filter_seqs(protein_dict):\n",
    "    re_chainbreak = re.compile(\"\\-*\\+*\\+\\-+\\+\\+*\\-*\")\n",
    "    keys_to_remove = []\n",
    "    for key in protein_dict.keys():\n",
    "        mask = protein_dict[key]['mask']\n",
    "        if re_chainbreak.search(mask):\n",
    "            keys_to_remove.append(key)\n",
    "        else:\n",
    "            coords = protein_dict[key]['tertiary']\n",
    "            coords_filt = filter_coords(coords, mask)\n",
    "            protein_dict[key]['tertiary'] = coords_filt\n",
    "    for key in keys_to_remove:\n",
    "        print(key, \" has a chainbreak, removing...\")\n",
    "        del protein_dict[key]\n",
    "\n",
    "def new_dihedral(p0, p1, p2, p3): # COPY PASTA'D FROM STACKEXCHANGE\n",
    "    \"\"\"Praxeolitic formula\n",
    "    1 sqrt, 1 cross product\"\"\"\n",
    "\n",
    "    b0 = -1.0*(p1 - p0)\n",
    "    b1 = p2 - p1\n",
    "    b2 = p3 - p2\n",
    "\n",
    "    # normalize b1 so that it does not influence magnitude of vector\n",
    "    # rejections that come next\n",
    "    b1 /= np.linalg.norm(b1)\n",
    "\n",
    "    # vector rejections\n",
    "    # v = projection of b0 onto plane perpendicular to b1\n",
    "    #   = b0 minus component that aligns with b1\n",
    "    # w = projection of b2 onto plane perpendicular to b1\n",
    "    #   = b2 minus component that aligns with b1\n",
    "    v = b0 - np.dot(b0, b1)*b1\n",
    "    w = b2 - np.dot(b2, b1)*b1\n",
    "\n",
    "    # angle between v and w in a plane is the torsion angle\n",
    "    # v and w may not be normalized but that's fine since tan is y/x\n",
    "    x = np.dot(v, w)\n",
    "    y = np.dot(np.cross(b1, v), w)\n",
    "    #return np.degrees(np.arctan2(y, x))\n",
    "    return np.arctan2(y, x)\n",
    "\n",
    "def calc_angles(coords):\n",
    "    N = coords.shape[1]\n",
    "    angles_all = []\n",
    "    for i in range(0, N-3, 3):\n",
    "        psi = new_dihedral(*[coords[:, x] for x in range(i, i+4)])\n",
    "        omega = new_dihedral(*[coords[:, x] for x in range(i+1, i+5)])\n",
    "        phi = new_dihedral(*[coords[:, x] for x in range(i+2, i+6)])\n",
    "        angles = [psi, omega, phi]\n",
    "        angles_all.append(angles)\n",
    "    return(np.array(angles_all))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "#for key in seqs.keys():\n",
    "#    print(key, seqs[key]['mask'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TBM#T0366  has a chainbreak, removing...\n",
      "TBM#T0365  has a chainbreak, removing...\n",
      "TBM#T0364  has a chainbreak, removing...\n",
      "TBM-hard#T0356  has a chainbreak, removing...\n",
      "TBM#T0339  has a chainbreak, removing...\n",
      "TBM#T0359  has a chainbreak, removing...\n",
      "TBM#T0331  has a chainbreak, removing...\n",
      "TBM#T0320  has a chainbreak, removing...\n",
      "TBM-hard#T0321  has a chainbreak, removing...\n",
      "TBM#T0313  has a chainbreak, removing...\n",
      "TBM#T0330  has a chainbreak, removing...\n",
      "TBM#T0301  has a chainbreak, removing...\n",
      "TBM#T0379  has a chainbreak, removing...\n",
      "TBM#T0305  has a chainbreak, removing...\n",
      "TBM#T0378  has a chainbreak, removing...\n",
      "TBM#T0288  has a chainbreak, removing...\n",
      "FM#T0296  has a chainbreak, removing...\n",
      "TBM-hard#T0316  has a chainbreak, removing...\n",
      "TBM#T0292  has a chainbreak, removing...\n",
      "TBM#T0293  has a chainbreak, removing...\n",
      "TBM#T0341  has a chainbreak, removing...\n",
      "TBM#T0291  has a chainbreak, removing...\n",
      "FM#T0300  has a chainbreak, removing...\n",
      "TBM-hard#T0347  has a chainbreak, removing...\n",
      "TBM#T0373  has a chainbreak, removing...\n",
      "TBM#T0384  has a chainbreak, removing...\n",
      "TBM#T0326  has a chainbreak, removing...\n"
     ]
    }
   ],
   "source": [
    "filter_seqs(seqs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "66"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(seqs.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(3, 297)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "seqs['TBM#T0285']['tertiary'].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "omega ca to ca, phi c to c, psi n to n, loop over every residue except first (or last) and calculate dihedral angle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_coords = seqs['TBM#T0285']['tertiary']\n",
    "test_angles = []\n",
    "for i in range(0, 297-3, 3):\n",
    "    psi = new_dihedral(*[test_coords[:, x] for x in range(i, i+4)])\n",
    "    omega = new_dihedral(*[test_coords[:, x] for x in range(i+1, i+5)])\n",
    "    phi = new_dihedral(*[test_coords[:, x] for x in range(i+2, i+6)])\n",
    "    angles = [psi, omega, phi]\n",
    "    test_angles.append(angles)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-0.7731468104941539"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_dihedral(*[test_coords[:, x] for x in range(0,0+4)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(98, 3)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.array(test_angles).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "530"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# check how big the proteins are, dont wanna pad a million zeroes\n",
    "keylengths=[len(seqs[key]['primary']) for key in seqs.keys()]\n",
    "max(keylengths)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_length = 1000\n",
    "sequence_list = []\n",
    "angle_list = []\n",
    "for key in seqs.keys():\n",
    "    sequence = aa_to_onehot(seqs[key]['primary'], aa_id_dict, seqs[key]['mask'])\n",
    "    coords = seqs[key]['tertiary']\n",
    "    angles = calc_angles(coords).T\n",
    "    if max_length-sequence.shape[0] < 0:\n",
    "        print(key, ' exceeds max length: ', max_length, ' , skipping...')\n",
    "        continue\n",
    "    else:\n",
    "        sequence_padded = np.pad(sequence, pad_width=((0,0), (0, max_length-sequence.shape[1])), constant_values=0)\n",
    "        sequence_list.append(sequence_padded)\n",
    "        angles_padded = np.pad(angles, pad_width=((0,0), (0, max_length-sequence.shape[1])), constant_values=0)\n",
    "        angle_list.append(angles_padded)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(66, 20, 1000)"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.array(sequence_list).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(66, 3, 999)"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.array(angle_list).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3.1415832609565535"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.max(np.array(angle_list))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "99.0"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.pad(aa_to_onehot(seqs['TBM#T0285']['primary'], aa_id_dict, seqs['TBM#T0285']['mask']),\n",
    "       pad_width=((0,0), (0,1000)), constant_values=0 ).sum() # making sure nothing weird goes on"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "#for i,j in zip(sequence_list, [seqs[key]['tertiary'].shape[1]/3 for key in seqs.keys()]):\n",
    "#    print(i.sum(), j) # sanity check"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def split_inds(N, train_fraction, validation_fraction):\n",
    "    indices = np.arange(N)\n",
    "    train_to = int(np.floor(len(indices)*train_fraction))\n",
    "    train_inds = indices[0: train_to]\n",
    "    val_to = int(np.floor(len(indices)*(train_fraction+validation_fraction)))\n",
    "    val_inds = indices[train_to: val_to]\n",
    "    test_inds = indices[val_to:]\n",
    "    return(train_inds, val_inds, test_inds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([66, 1, 20, 1000])\n"
     ]
    }
   ],
   "source": [
    "sequence_tensor = torch.tensor(sequence_list).float()\n",
    "sequence_tensor = sequence_tensor.unsqueeze(1)\n",
    "print(sequence_tensor.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([66, 1, 3, 999])\n"
     ]
    }
   ],
   "source": [
    "angle_tensor = torch.tensor(angle_list).float()\n",
    "angle_tensor = angle_tensor.unsqueeze(1)\n",
    "print(angle_tensor.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.float32"
      ]
     },
     "execution_count": 112,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sequence_tensor.dtype"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_inds, val_inds, test_inds = split_inds(66, 0.8, 0.1)\n",
    "\n",
    "train_data = [sequence_tensor[train_inds, :, :, :], angle_tensor[train_inds, :, :, :]]\n",
    "val_data = [sequence_tensor[val_inds, :, :, :], angle_tensor[val_inds, :, :, :]]\n",
    "test_data = [sequence_tensor[test_inds, :, :, :], angle_tensor[test_inds, :, :, :]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import Dataset\n",
    "class proteindataset(Dataset):\n",
    "    def __init__(self, seqs, angles):\n",
    "        self.sequences = seqs\n",
    "        self.angles = angles\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.sequences.shape[0]\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return [self.sequences[idx, :, :, :], self.angles[idx, :, :, :]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = proteindataset(train_data[0], train_data[1])\n",
    "val_dataset = proteindataset(val_data[0], val_data[1])\n",
    "test_dataset = proteindataset(test_data[0], test_data[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 116,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainloader = torch.utils.data.DataLoader(train_dataset, batch_size=4,\n",
    "                                          shuffle=True, num_workers=2)\n",
    "valloader = torch.utils.data.DataLoader(val_dataset, batch_size=4,\n",
    "                                         shuffle=False, num_workers=2)\n",
    "testloader = torch.utils.data.DataLoader(test_dataset, batch_size=4,\n",
    "                                         shuffle=False, num_workers=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 122,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "\n",
    "\n",
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 6, (20, 4)) #in, out, kernel size (can be 1 number)\n",
    "        self.conv2 = nn.Conv2d(6, 12, (1, 4))\n",
    "        self.deconv1 = nn.ConvTranspose2d(in_channels=12, out_channels=6, kernel_size=(1, 4))\n",
    "        self.deconv2 = nn.ConvTranspose2d(in_channels=6, out_channels=3, kernel_size=(1, 3))\n",
    "\n",
    "    def forward(self, x):\n",
    "        #print('input shape: ', x.shape)\n",
    "        conv1_out = F.relu(self.conv1(x))\n",
    "        #print('conv1 out shape: ', conv1_out.shape)\n",
    "        conv2_out = F.relu(self.conv2(conv1_out))\n",
    "        #print('conv2 out shape: ', conv2_out.shape)\n",
    "        deconv1_out = F.relu(self.deconv1(conv2_out))\n",
    "        #print('deconv1 out shape: ', deconv1_out.shape)\n",
    "        deconv2_out = self.deconv2(deconv1_out)\n",
    "        #print('deconv2 out shape: ', deconv2_out.shape)\n",
    "        return deconv2_out\n",
    "\n",
    "\n",
    "net = Net()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 158,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4.0\n",
      "epoch: 0, iteration: 3] loss: 0.8308959752321243\n",
      "0.8133204430341721 0.8133204430341721\n",
      "new best loss, saving..\n",
      "epoch: 0, iteration: 7] loss: 0.8133204430341721\n",
      "0.8099523782730103 0.8099523782730103\n",
      "new best loss, saving..\n",
      "epoch: 0, iteration: 11] loss: 0.8099523782730103\n",
      "0.7754370272159576 0.7754370272159576\n",
      "new best loss, saving..\n",
      "epoch: 1, iteration: 3] loss: 0.7754370272159576\n",
      "0.7317431718111038 0.7317431718111038\n",
      "new best loss, saving..\n",
      "epoch: 1, iteration: 7] loss: 0.7317431718111038\n",
      "0.9498272091150284 0.7317431718111038\n",
      "epoch: 1, iteration: 11] loss: 0.9498272091150284\n",
      "1.0188564360141754 0.7317431718111038\n",
      "epoch: 2, iteration: 3] loss: 1.0188564360141754\n",
      "0.7587235197424889 0.7317431718111038\n",
      "epoch: 2, iteration: 7] loss: 0.7587235197424889\n",
      "0.6247512698173523 0.6247512698173523\n",
      "new best loss, saving..\n",
      "epoch: 2, iteration: 11] loss: 0.6247512698173523\n",
      "0.8720566481351852 0.6247512698173523\n",
      "epoch: 3, iteration: 3] loss: 0.8720566481351852\n",
      "0.8232081383466721 0.6247512698173523\n",
      "epoch: 3, iteration: 7] loss: 0.8232081383466721\n",
      "0.7570068538188934 0.6247512698173523\n",
      "epoch: 3, iteration: 11] loss: 0.7570068538188934\n",
      "0.8684588298201561 0.6247512698173523\n",
      "epoch: 4, iteration: 3] loss: 0.8684588298201561\n",
      "0.8021431565284729 0.6247512698173523\n",
      "epoch: 4, iteration: 7] loss: 0.8021431565284729\n",
      "0.710144504904747 0.6247512698173523\n",
      "epoch: 4, iteration: 11] loss: 0.710144504904747\n",
      "0.8333712667226791 0.6247512698173523\n",
      "epoch: 5, iteration: 3] loss: 0.8333712667226791\n",
      "0.7880004197359085 0.6247512698173523\n",
      "epoch: 5, iteration: 7] loss: 0.7880004197359085\n",
      "0.8351570516824722 0.6247512698173523\n",
      "epoch: 5, iteration: 11] loss: 0.8351570516824722\n",
      "0.7789084613323212 0.6247512698173523\n",
      "epoch: 6, iteration: 3] loss: 0.7789084613323212\n",
      "0.7976109534502029 0.6247512698173523\n",
      "epoch: 6, iteration: 7] loss: 0.7976109534502029\n",
      "0.8008138239383698 0.6247512698173523\n",
      "epoch: 6, iteration: 11] loss: 0.8008138239383698\n",
      "0.8161664009094238 0.6247512698173523\n",
      "epoch: 7, iteration: 3] loss: 0.8161664009094238\n",
      "0.930717870593071 0.6247512698173523\n",
      "epoch: 7, iteration: 7] loss: 0.930717870593071\n",
      "0.7248279303312302 0.6247512698173523\n",
      "epoch: 7, iteration: 11] loss: 0.7248279303312302\n",
      "0.7865548506379128 0.6247512698173523\n",
      "epoch: 8, iteration: 3] loss: 0.7865548506379128\n",
      "0.7884890884160995 0.6247512698173523\n",
      "epoch: 8, iteration: 7] loss: 0.7884890884160995\n",
      "0.8671759068965912 0.6247512698173523\n",
      "epoch: 8, iteration: 11] loss: 0.8671759068965912\n",
      "0.7858400940895081 0.6247512698173523\n",
      "epoch: 9, iteration: 3] loss: 0.7858400940895081\n",
      "0.7077689617872238 0.6247512698173523\n",
      "epoch: 9, iteration: 7] loss: 0.7077689617872238\n",
      "0.923159584403038 0.6247512698173523\n",
      "epoch: 9, iteration: 11] loss: 0.923159584403038\n",
      "Finished Training\n"
     ]
    }
   ],
   "source": [
    "criterion = nn.MSELoss()\n",
    "optimizer = optim.Adam(net.parameters(), lr=1e-4)\n",
    "\n",
    "prints_per_epoch = 3\n",
    "\n",
    "verbose_k = np.floor(len(trainloader)/prints_per_epoch)\n",
    "print(verbose_k)\n",
    "\n",
    "losses = []\n",
    "iterations = []\n",
    "best_loss = None\n",
    "\n",
    "for epoch in range(10):\n",
    "\n",
    "    running_loss = 0.0\n",
    "    for i, data in enumerate(trainloader):\n",
    "    #for i in range(66):\n",
    "        #sequence = train_data[0][i, :, :, :].unsqueeze(0)\n",
    "        #true_angles = train_data[1][i, :, :, :].unsqueeze(0)\n",
    "        sequence, true_angles = data\n",
    "        #print(sequence.shape, true_angles.shape)\n",
    "\n",
    "        # zero the parameter gradients\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # forward + backward + optimize\n",
    "        predicted_angles = net(sequence)\n",
    "\n",
    "        loss = criterion(predicted_angles, true_angles)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        # print statistics, should add validation loss\n",
    "        running_loss += loss.item()\n",
    "        \n",
    "        if (i+1) % verbose_k == 0:   \n",
    "            losses.append(running_loss/verbose_k)\n",
    "            true_iter = len(trainloader)*epoch + i\n",
    "            iterations.append(true_iter)\n",
    "            \n",
    "            if best_loss == None:\n",
    "                best_loss = running_loss/verbose_k\n",
    "            else:\n",
    "                if running_loss/verbose_k <= min(losses):\n",
    "                    print('new best loss, saving..')\n",
    "                    best_loss = running_loss/verbose_k\n",
    "                    torch.save(net.state_dict(), 'best_fcn_parameters.pt')\n",
    "            \n",
    "            print('epoch: {}, iteration: {}] loss: {}'.format(epoch, i, running_loss/verbose_k))\n",
    "            running_loss = 0.0\n",
    "\n",
    "print('Finished Training')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 144,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "13"
      ]
     },
     "execution_count": 144,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(trainloader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 145,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2]"
      ]
     },
     "execution_count": 145,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[len(x) for x in trainloader]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 146,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[torch.Size([52, 1, 20, 1000]), torch.Size([52, 1, 3, 999])]"
      ]
     },
     "execution_count": 146,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[x.shape for x in train_data]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 147,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 3, 999])"
      ]
     },
     "execution_count": 147,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_data[1][0, :, : ,:].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 159,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.6247512698173523"
      ]
     },
     "execution_count": 159,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "best_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 160,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.6247512698173523"
      ]
     },
     "execution_count": 160,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "min(losses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}