_2_VAE_train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from imp import reload\n",
    "import h5py\n",
    "import shutil\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import init\n",
    "from torch.autograd import Variable\n",
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "from itertools import chain\n",
    "from sklearn.model_selection import KFold\n",
    "from progiter import ProgIter\n",
    "from sklearn.model_selection import train_test_split\n",
    "from mymodule import VAE_model\n",
    "\n",
    "random_state=34"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<module 'mymodule.VAE_model' from '/tf/notebooks/code_for_pub/mymodule/VAE_model.py'>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "reload(VAE_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# all drugs\n",
    "with open('/tf/notebooks/code_for_pub/smiles_files/smiles_drugcombANDchembl26.pickle','rb') as f:\n",
    "    a = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('/tf/notebooks/code_for_pub/smiles_files/smiles_drugcomb_BY_cid_duplicated.pickle','rb') as f:\n",
    "    b = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1795483 unique SMILES read in 11.0767 seconds\n"
     ]
    }
   ],
   "source": [
    "processor = VAE_model.ProcessSMILES(smiles_chembl=a, smiles_dc=b)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 256 bit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "one-hot-encoded 1795483 SMILES in 225.969 seconds\n"
     ]
    }
   ],
   "source": [
    "# dataprep\n",
    "df = processor.one_hot_encoded(processor.smiles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'charset' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-8-f8dfb0938904>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      6\u001b[0m \u001b[0;31m#best = torch.load('/tf/notebooks/code_for_pub/_logs_as_python_files/vae_training_logs/model_best.pth.tar')\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 7\u001b[0;31m \u001b[0mencoder\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mVAE_model\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mMolEncoder\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m140\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mo\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m256\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mc\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcharset\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      8\u001b[0m \u001b[0;31m#encoder.load_state_dict(best['encoder'])\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      9\u001b[0m \u001b[0mencoder\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mapply\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mVAE_model\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minitialize_weights\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'charset' is not defined"
     ]
    }
   ],
   "source": [
    "batch_size = 1800\n",
    "patience = 3\n",
    "n_epochs = 10\n",
    "charset = processor.charset\n",
    "dtype=torch.cuda.FloatTensor\n",
    "\n",
    "#best = torch.load('/tf/notebooks/code_for_pub/_logs_as_python_files/vae_training_logs/model_best.pth.tar')\n",
    "encoder = VAE_model.MolEncoder(i=140, o=256, c=len(charset))\n",
    "#encoder.load_state_dict(best['encoder'])\n",
    "encoder.apply(VAE_model.initialize_weights)\n",
    "decoder = VAE_model.MolDecoder(i=140, o=256, c=len(charset))\n",
    "decoder.apply(VAE_model.initialize_weights)\n",
    "#decoder.load_state_dict(best['decoder'])\n",
    "encoder.cuda()\n",
    "decoder.cuda()\n",
    "best_loss = 1e6\n",
    "#best_loss = best['avg_val_loss']\n",
    "\n",
    "optimizer = torch.optim.Adam(chain(encoder.parameters(), decoder.parameters()), lr=1e-3)\n",
    "scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, \n",
    "                                                       factor=0.2,\n",
    "                                                       patience=patience,\n",
    "                                                       mode='min', \n",
    "                                                       min_lr=1e-6)\n",
    "\n",
    "kf = KFold(n_splits=5, shuffle=True, random_state=random_state)        \n",
    "for ind,(train_index, test_index) in enumerate(kf.split(df)):\n",
    "    \n",
    "    # make loader train\n",
    "    dt = torch.from_numpy(df[train_index])\n",
    "    train = TensorDataset(dt, torch.zeros(dt.size()[0]))\n",
    "    train_loader = DataLoader(train, batch_size=batch_size, shuffle=True)\n",
    "\n",
    "    # make loader test\n",
    "    dte = torch.from_numpy(df[test_index])\n",
    "    test = TensorDataset(dte, torch.zeros(dte.size()[0]))\n",
    "    val_loader = DataLoader(test, batch_size=batch_size, shuffle=True)\n",
    "    del dt, dte, train, test\n",
    "    \n",
    "    # for other folds we re-use previous best model's weights\n",
    "    if ind != 0:\n",
    "        encoder = VAE_model.MolEncoder(i=140, o=256, c=len(charset))\n",
    "        best = torch.load('/tf/notebooks/code_for_pub/_logs_as_python_files/vae_training_logs/model_best.pth.tar')\n",
    "        encoder.load_state_dict(best['encoder'])\n",
    "        decoder = VAE_model.MolDecoder(i=140, o=256, c=len(charset))\n",
    "        decoder.load_state_dict(best['decoder'])\n",
    "        encoder.cuda()\n",
    "        decoder.cuda()\n",
    "\n",
    "        optimizer = torch.optim.Adam(chain(encoder.parameters(), decoder.parameters()), lr=1e-3)\n",
    "        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, \n",
    "                                                               factor=0.1,\n",
    "                                                               patience=patience,\n",
    "                                                               mode='min', \n",
    "                                                               min_lr=1e-6a)\n",
    "    \n",
    "    \n",
    "    \n",
    "    # train for n epochs\n",
    "    for epoch in range(n_epochs):\n",
    "        print(f'outer fold:{ind}, epoch:{epoch}')\n",
    "        VAE_model.train_model(train_loader, encoder, decoder, optimizer, dtype)\n",
    "        avg_val_loss = VAE_model.validate_model(val_loader, encoder, decoder, dtype)\n",
    "\n",
    "        scheduler.step(avg_val_loss)\n",
    "\n",
    "        is_best = avg_val_loss < best_loss\n",
    "\n",
    "        if is_best:\n",
    "            best_loss = avg_val_loss\n",
    "        VAE_model.save_checkpoint({\n",
    "            'epoch': epoch,\n",
    "            'encoder': encoder.state_dict(),\n",
    "            'decoder': decoder.state_dict(),\n",
    "            'charset': charset,\n",
    "            'avg_val_loss': avg_val_loss,\n",
    "            'optimizer': optimizer.state_dict(),\n",
    "        }, is_best, size=256)\n",
    "\n",
    "    del train_loader, encoder, decoder, scheduler, optimizer \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 16 bit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1795483 unique SMILES read in 2.6206 seconds\n"
     ]
    }
   ],
   "source": [
    "processor = VAE_model.Process_SMILES()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "one-hot-encoded 1795483 SMILES in 253.2625 seconds\n"
     ]
    }
   ],
   "source": [
    "df= processor.one_hot_encode_all(processor.smiles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  0/10... rate=0 Hz, eta=?, total=0:00:00, wall=16:29 UTCouter fold:0, epoch:0\n",
      "t = 100, loss = 1.54854369\n",
      "t = 200, loss = 1.55469525\n",
      "t = 300, loss = 1.56053662\n",
      "t = 400, loss = 1.50880802\n",
      "t = 500, loss = 1.53006780\n",
      "t = 600, loss = 1.54756558\n",
      "t = 700, loss = 1.54510748\n",
      "average validation loss: 1.52949405\n",
      "  1/10... rate=0.00 Hz, eta=1:53:39, total=0:12:37, wall=16:42 UTCouter fold:0, epoch:1\n",
      "t = 100, loss = 1.54035807\n",
      "t = 200, loss = 1.49055052\n",
      "t = 300, loss = 1.53780365\n",
      "t = 400, loss = 1.53209603\n",
      "t = 500, loss = 1.79145467\n",
      "t = 600, loss = 1.58308709\n",
      "t = 700, loss = 1.49396610\n",
      "average validation loss: 1.51151991\n",
      "  2/10... rate=0.00 Hz, eta=1:40:57, total=0:25:14, wall=16:55 UTCouter fold:0, epoch:2\n",
      "t = 100, loss = 1.49950171\n",
      "t = 200, loss = 1.47682941\n",
      "t = 300, loss = 1.50885391\n",
      "t = 400, loss = 1.51407170\n",
      "t = 500, loss = 1.51095748\n",
      "t = 600, loss = 1.49812484\n",
      "t = 700, loss = 1.58757770\n",
      "average validation loss: 1.48061466\n",
      "  3/10... rate=0.00 Hz, eta=1:28:17, total=0:37:50, wall=17:07 UTCouter fold:0, epoch:3\n",
      "t = 100, loss = 1.41788411\n",
      "t = 200, loss = 1.55312359\n",
      "t = 300, loss = 1.45899451\n",
      "t = 400, loss = 1.42670798\n",
      "t = 500, loss = 1.43479776\n",
      "t = 600, loss = 1.43175042\n",
      "t = 700, loss = 1.43048298\n",
      "average validation loss: 1.48656344\n",
      "  4/10... rate=0.00 Hz, eta=1:15:39, total=0:50:26, wall=17:20 UTCouter fold:0, epoch:4\n",
      "t = 100, loss = 1.42927539\n",
      "t = 200, loss = 1.46906590\n",
      "t = 300, loss = 1.42599761\n",
      "t = 400, loss = 1.43950140\n",
      "t = 500, loss = 1.38266993\n",
      "t = 600, loss = 1.43017173\n",
      "t = 700, loss = 1.46615064\n",
      "average validation loss: 1.48452914\n",
      "  5/10... rate=0.00 Hz, eta=1:03:02, total=1:03:02, wall=17:32 UTCouter fold:0, epoch:5\n",
      "t = 100, loss = 1.41227067\n",
      "t = 200, loss = 1.43223536\n",
      "t = 300, loss = 1.46205962\n",
      "t = 400, loss = 1.37543023\n",
      "t = 500, loss = 1.39844954\n",
      "t = 600, loss = 1.40275252\n",
      "t = 700, loss = 1.38560891\n",
      "average validation loss: 1.39508700\n",
      "  6/10... rate=0.00 Hz, eta=0:50:25, total=1:15:38, wall=17:45 UTCouter fold:0, epoch:6\n",
      "t = 100, loss = 1.39348006\n",
      "t = 200, loss = 1.36579943\n",
      "t = 300, loss = 1.34645176\n",
      "t = 400, loss = 1.37400985\n",
      "t = 500, loss = 1.35966587\n",
      "t = 600, loss = 1.46034241\n",
      "t = 700, loss = 1.38935220\n",
      "average validation loss: 1.38620424\n",
      "  7/10... rate=0.00 Hz, eta=0:37:49, total=1:28:14, wall=17:58 UTCouter fold:0, epoch:7\n",
      "t = 100, loss = 1.37270761\n",
      "t = 200, loss = 1.36950600\n",
      "t = 300, loss = 1.35851061\n",
      "t = 400, loss = 1.34850442\n",
      "t = 500, loss = 1.35648572\n",
      "t = 600, loss = 1.36410272\n",
      "t = 700, loss = 1.28811622\n",
      "average validation loss: 1.34213209\n",
      "  8/10... rate=0.00 Hz, eta=0:25:12, total=1:40:51, wall=18:10 UTCouter fold:0, epoch:8\n",
      "t = 100, loss = 1.31881571\n",
      "t = 200, loss = 1.38283944\n",
      "t = 300, loss = 1.28422177\n",
      "t = 400, loss = 1.29363608\n",
      "t = 500, loss = 1.31277466\n",
      "t = 600, loss = 1.29841363\n",
      "t = 700, loss = 1.31462085\n",
      "average validation loss: 1.35326862\n",
      "  9/10... rate=0.00 Hz, eta=0:12:36, total=1:53:27, wall=18:23 UTCouter fold:0, epoch:9\n",
      "t = 100, loss = 1.32450163\n",
      "t = 200, loss = 1.32717252\n",
      "t = 300, loss = 1.36184931\n",
      "t = 400, loss = 1.34653771\n",
      "t = 500, loss = 1.34057403\n",
      "t = 600, loss = 1.33351326\n",
      "t = 700, loss = 1.30686867\n",
      "average validation loss: 1.31991529\n",
      " 10/10... rate=0.00 Hz, eta=0:00:00, total=2:06:03, wall=18:35 UTC\n",
      "  0/10... rate=0 Hz, eta=?, total=0:00:00, wall=18:37 UTCouter fold:1, epoch:0\n",
      "t = 100, loss = 1.29662287\n",
      "t = 200, loss = 1.33422852\n",
      "t = 300, loss = 1.33458328\n",
      "t = 400, loss = 1.39335024\n",
      "t = 500, loss = 1.30674326\n",
      "t = 600, loss = 1.29679871\n",
      "t = 700, loss = 1.33070660\n",
      "average validation loss: 1.29847848\n",
      "  1/10... rate=0.00 Hz, eta=1:53:30, total=0:12:36, wall=18:49 UTCouter fold:1, epoch:1\n",
      "t = 100, loss = 1.34102559\n",
      "t = 200, loss = 1.28115153\n",
      "t = 300, loss = 1.27709019\n",
      "t = 400, loss = 1.29666376\n",
      "t = 500, loss = 1.25557244\n",
      "t = 600, loss = 1.32040417\n",
      "t = 700, loss = 1.30409741\n",
      "average validation loss: 1.33777082\n",
      "  2/10... rate=0.00 Hz, eta=1:40:49, total=0:25:12, wall=19:02 UTCouter fold:1, epoch:2\n",
      "t = 100, loss = 1.28804898\n",
      "t = 200, loss = 1.26813149\n",
      "t = 300, loss = 1.28268194\n",
      "t = 400, loss = 1.31586301\n",
      "t = 500, loss = 1.27907395\n",
      "t = 600, loss = 1.29040515\n",
      "t = 700, loss = 1.24255013\n",
      "average validation loss: 1.26233983\n",
      "  3/10... rate=0.00 Hz, eta=1:28:17, total=0:37:50, wall=19:15 UTCouter fold:1, epoch:3\n",
      "t = 100, loss = 1.35895348\n",
      "t = 200, loss = 1.25642407\n",
      "t = 300, loss = 1.34132028\n",
      "t = 400, loss = 1.20125020\n",
      "t = 500, loss = 1.24092448\n",
      "t = 600, loss = 1.23371184\n",
      "t = 700, loss = 1.27799165\n",
      "average validation loss: 1.26372921\n",
      "  4/10... rate=0.00 Hz, eta=1:15:40, total=0:50:26, wall=19:27 UTCouter fold:1, epoch:4\n",
      "t = 100, loss = 1.25053549\n",
      "t = 200, loss = 1.18109834\n",
      "t = 300, loss = 1.24626851\n",
      "t = 400, loss = 1.20785046\n",
      "t = 500, loss = 1.21525204\n",
      "t = 600, loss = 1.24493182\n",
      "t = 700, loss = 1.19910872\n",
      "average validation loss: 1.30045485\n",
      "  5/10... rate=0.00 Hz, eta=1:03:02, total=1:03:02, wall=19:40 UTCouter fold:1, epoch:5\n",
      "t = 100, loss = 1.22544634\n",
      "t = 200, loss = 1.20517886\n",
      "t = 300, loss = 1.19620860\n",
      "t = 400, loss = 1.21267927\n",
      "t = 500, loss = 1.21304250\n",
      "t = 600, loss = 1.23587298\n",
      "t = 700, loss = 1.18474078\n",
      "average validation loss: 1.24043548\n",
      "  6/10... rate=0.00 Hz, eta=0:50:25, total=1:15:38, wall=19:52 UTCouter fold:1, epoch:6\n",
      "t = 100, loss = 1.10047209\n",
      "t = 200, loss = 1.19347787\n",
      "t = 300, loss = 1.18405783\n",
      "t = 400, loss = 1.23960340\n",
      "t = 500, loss = 1.20602274\n",
      "t = 600, loss = 1.22491717\n",
      "t = 700, loss = 1.18958259\n",
      "average validation loss: 1.20135450\n",
      "  7/10... rate=0.00 Hz, eta=0:37:49, total=1:28:14, wall=20:05 UTCouter fold:1, epoch:7\n",
      "t = 100, loss = 1.16371942\n",
      "t = 200, loss = 1.25462008\n",
      "t = 300, loss = 1.26314867\n",
      "t = 400, loss = 1.18716311\n",
      "t = 500, loss = 1.32829368\n",
      "t = 600, loss = 1.24309003\n",
      "t = 700, loss = 1.15356362\n",
      "average validation loss: 1.20361626\n",
      "  8/10... rate=0.00 Hz, eta=0:25:12, total=1:40:50, wall=20:18 UTCouter fold:1, epoch:8\n",
      "t = 100, loss = 1.21013355\n",
      "t = 200, loss = 1.12712574\n",
      "t = 300, loss = 1.18853974\n",
      "t = 400, loss = 1.22226441\n",
      "t = 500, loss = 1.14982247\n",
      "t = 600, loss = 1.16133952\n",
      "t = 700, loss = 1.16366816\n",
      "average validation loss: 1.16843593\n",
      "  9/10... rate=0.00 Hz, eta=0:12:36, total=1:53:27, wall=20:30 UTCouter fold:1, epoch:9\n",
      "t = 100, loss = 1.22465467\n",
      "t = 200, loss = 1.15910053\n",
      "t = 300, loss = 1.11580729\n",
      "t = 400, loss = 1.15743792\n",
      "t = 500, loss = 1.15818846\n",
      "t = 600, loss = 1.13701105\n",
      "t = 700, loss = 1.15955067\n",
      "average validation loss: 1.17464435\n",
      " 10/10... rate=0.00 Hz, eta=0:00:00, total=2:06:02, wall=20:43 UTC\n",
      "  0/10... rate=0 Hz, eta=?, total=0:00:00, wall=20:43 UTCouter fold:2, epoch:0\n",
      "t = 100, loss = 1.13630295\n",
      "t = 200, loss = 1.15097082\n",
      "t = 300, loss = 1.11885154\n",
      "t = 400, loss = 1.21201909\n",
      "t = 500, loss = 1.12459648\n",
      "t = 600, loss = 1.15793419\n",
      "t = 700, loss = 1.19693398\n",
      "average validation loss: 1.16614068\n",
      "  1/10... rate=0.00 Hz, eta=1:53:20, total=0:12:35, wall=20:56 UTCouter fold:2, epoch:1\n",
      "t = 100, loss = 1.15386450\n",
      "t = 200, loss = 1.11642241\n",
      "t = 300, loss = 1.13997686\n",
      "t = 400, loss = 1.13128829\n",
      "t = 500, loss = 1.11615992\n",
      "t = 600, loss = 1.16566634\n",
      "t = 700, loss = 1.12640846\n",
      "average validation loss: 1.14072728\n",
      "  2/10... rate=0.00 Hz, eta=1:40:48, total=0:25:12, wall=21:08 UTCouter fold:2, epoch:2\n",
      "t = 100, loss = 1.17964947\n",
      "t = 200, loss = 1.14328563\n",
      "t = 300, loss = 1.13559890\n",
      "t = 400, loss = 1.15011656\n",
      "t = 500, loss = 1.15708268\n",
      "t = 600, loss = 1.13270628\n",
      "t = 700, loss = 1.20044911\n",
      "average validation loss: 1.14194763\n",
      "  3/10... rate=0.00 Hz, eta=1:28:11, total=0:37:47, wall=21:21 UTCouter fold:2, epoch:3\n",
      "t = 100, loss = 1.10556114\n",
      "t = 200, loss = 1.15119982\n",
      "t = 300, loss = 1.17715263\n",
      "t = 400, loss = 1.17010880\n",
      "t = 500, loss = 1.17321205\n",
      "t = 600, loss = 1.12477422\n",
      "t = 700, loss = 1.13873506\n",
      "average validation loss: 1.14145887\n",
      "  4/10... rate=0.00 Hz, eta=1:15:34, total=0:50:23, wall=21:34 UTCouter fold:2, epoch:4\n",
      "t = 100, loss = 1.13801575\n",
      "t = 200, loss = 1.09111249\n",
      "t = 300, loss = 1.12079597\n",
      "t = 400, loss = 1.08845532\n",
      "t = 500, loss = 1.11513269\n",
      "t = 600, loss = 1.11181629\n",
      "t = 700, loss = 1.08935153\n",
      "average validation loss: 1.14174318\n",
      "  5/10... rate=0.00 Hz, eta=1:02:59, total=1:02:59, wall=21:46 UTCouter fold:2, epoch:5\n",
      "t = 100, loss = 1.13863873\n",
      "t = 200, loss = 1.08962035\n",
      "t = 300, loss = 1.04693449\n",
      "t = 400, loss = 1.13216150\n",
      "t = 500, loss = 1.13957143\n",
      "t = 600, loss = 1.18491709\n",
      "t = 700, loss = 1.09795332\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "average validation loss: 1.11699271\n",
      "  6/10... rate=0.00 Hz, eta=0:50:23, total=1:15:35, wall=21:59 UTCouter fold:2, epoch:6\n",
      "t = 100, loss = 1.11760032\n",
      "t = 200, loss = 1.10683370\n",
      "t = 300, loss = 1.10707080\n",
      "t = 400, loss = 1.12318027\n",
      "t = 500, loss = 1.08443213\n",
      "t = 600, loss = 1.03949726\n",
      "t = 700, loss = 1.15147984\n",
      "average validation loss: 1.11728251\n",
      "  7/10... rate=0.00 Hz, eta=0:37:47, total=1:28:11, wall=22:11 UTCouter fold:2, epoch:7\n",
      "t = 100, loss = 1.06308961\n",
      "t = 200, loss = 1.08312571\n",
      "t = 300, loss = 1.07712138\n",
      "t = 400, loss = 1.04331720\n",
      "t = 500, loss = 1.05978858\n",
      "t = 600, loss = 1.13209796\n",
      "t = 700, loss = 1.09352243\n",
      "average validation loss: 1.11491477\n",
      "  8/10... rate=0.00 Hz, eta=0:25:11, total=1:40:47, wall=22:24 UTCouter fold:2, epoch:8\n",
      "t = 100, loss = 1.14065921\n",
      "t = 200, loss = 1.15282583\n",
      "t = 300, loss = 1.09161639\n",
      "t = 400, loss = 1.12661481\n",
      "t = 500, loss = 1.14359820\n",
      "t = 600, loss = 1.16311002\n",
      "t = 700, loss = 1.07669079\n",
      "average validation loss: 1.10470498\n",
      "  9/10... rate=0.00 Hz, eta=0:12:35, total=1:53:23, wall=22:37 UTCouter fold:2, epoch:9\n",
      "t = 100, loss = 1.07078183\n",
      "t = 200, loss = 1.10677683\n",
      "t = 300, loss = 1.05789244\n",
      "t = 400, loss = 1.06888604\n",
      "t = 500, loss = 1.09976161\n",
      "t = 600, loss = 1.10510552\n",
      "t = 700, loss = 1.08568048\n",
      "average validation loss: 1.09810257\n",
      " 10/10... rate=0.00 Hz, eta=0:00:00, total=2:05:59, wall=22:49 UTC\n",
      "  0/10... rate=0 Hz, eta=?, total=0:00:00, wall=22:50 UTCouter fold:3, epoch:0\n",
      "t = 100, loss = 2.86160326\n",
      "t = 200, loss = 2.23041725\n",
      "t = 300, loss = 1.77549708\n",
      "t = 400, loss = 1.48834109\n",
      "t = 500, loss = 1.25662601\n",
      "t = 600, loss = 1.19179356\n",
      "t = 700, loss = 1.13334560\n",
      "average validation loss: 1.12060416\n",
      "  1/10... rate=0.00 Hz, eta=1:53:31, total=0:12:36, wall=23:02 UTCouter fold:3, epoch:1\n",
      "t = 100, loss = 1.08221591\n",
      "t = 200, loss = 1.06347740\n",
      "t = 300, loss = 1.06245208\n",
      "t = 400, loss = 1.04656959\n",
      "t = 500, loss = 1.05802417\n",
      "t = 600, loss = 1.06424475\n",
      "t = 700, loss = 0.99644333\n",
      "average validation loss: 1.01893902\n",
      "  2/10... rate=0.00 Hz, eta=1:40:50, total=0:25:12, wall=23:15 UTCouter fold:3, epoch:2\n",
      "t = 100, loss = 0.99368107\n",
      "t = 200, loss = 0.99437416\n",
      "t = 300, loss = 1.06107235\n",
      "t = 400, loss = 0.96692115\n",
      "t = 500, loss = 1.02711272\n",
      "t = 600, loss = 0.99437708\n",
      "t = 700, loss = 0.95781112\n",
      "average validation loss: 1.01230621\n",
      "  3/10... rate=0.00 Hz, eta=1:28:15, total=0:37:49, wall=23:28 UTCouter fold:3, epoch:3\n",
      "t = 100, loss = 0.96416181\n",
      "t = 200, loss = 1.03398180\n",
      "t = 300, loss = 1.00419497\n",
      "t = 400, loss = 1.01551235\n",
      "t = 500, loss = 0.98697776\n",
      "t = 600, loss = 1.00791550\n",
      "t = 700, loss = 1.04248774\n",
      "average validation loss: 1.02136457\n",
      "  4/10... rate=0.00 Hz, eta=1:15:39, total=0:50:26, wall=23:40 UTCouter fold:3, epoch:4\n",
      "t = 100, loss = 0.98523849\n",
      "t = 200, loss = 0.97398704\n",
      "t = 300, loss = 1.01998460\n",
      "t = 400, loss = 0.98644185\n",
      "t = 500, loss = 1.02483022\n",
      "t = 600, loss = 1.05682266\n",
      "t = 700, loss = 1.07726920\n",
      "average validation loss: 1.00923359\n",
      "  5/10... rate=0.00 Hz, eta=1:03:01, total=1:03:01, wall=23:53 UTCouter fold:3, epoch:5\n",
      "t = 100, loss = 1.00520504\n",
      "t = 200, loss = 1.02386463\n",
      "t = 300, loss = 1.07237494\n",
      "t = 400, loss = 1.01941419\n",
      "t = 500, loss = 1.00114822\n",
      "t = 600, loss = 1.03400850\n",
      "t = 700, loss = 1.01397455\n",
      "average validation loss: 1.01221621\n",
      "  6/10... rate=0.00 Hz, eta=0:50:25, total=1:15:37, wall=00:06 UTCouter fold:3, epoch:6\n",
      "t = 100, loss = 0.98299557\n",
      "t = 200, loss = 0.99469364\n",
      "t = 300, loss = 1.04784155\n",
      "t = 400, loss = 1.00816321\n",
      "t = 500, loss = 0.98269355\n",
      "t = 600, loss = 1.03995252\n",
      "t = 700, loss = 1.04303145\n",
      "average validation loss: 1.03693998\n",
      "  7/10... rate=0.00 Hz, eta=0:37:48, total=1:28:14, wall=00:18 UTCouter fold:3, epoch:7\n",
      "t = 100, loss = 1.06521714\n",
      "t = 200, loss = 1.04788733\n",
      "t = 300, loss = 0.98098326\n",
      "t = 400, loss = 1.00910771\n",
      "t = 500, loss = 1.03867686\n",
      "t = 600, loss = 1.06886554\n",
      "t = 700, loss = 1.05465961\n",
      "average validation loss: 1.04192221\n",
      "  8/10... rate=0.00 Hz, eta=0:25:12, total=1:40:49, wall=00:31 UTCouter fold:3, epoch:8\n",
      "t = 100, loss = 1.01110339\n",
      "t = 200, loss = 1.05285859\n",
      "t = 300, loss = 1.02967417\n",
      "t = 400, loss = 1.04905307\n",
      "t = 500, loss = 1.02643466\n",
      "t = 600, loss = 1.07965326\n",
      "t = 700, loss = 0.99793285\n",
      "average validation loss: 1.03690028\n",
      "  9/10... rate=0.00 Hz, eta=0:12:36, total=1:53:25, wall=00:43 UTCouter fold:3, epoch:9\n",
      "t = 100, loss = 0.93660855\n",
      "t = 200, loss = 0.94276279\n",
      "t = 300, loss = 0.93465334\n",
      "t = 400, loss = 0.91056329\n",
      "t = 500, loss = 0.97506386\n",
      "t = 600, loss = 0.93649012\n",
      "t = 700, loss = 0.93178946\n",
      "average validation loss: 0.95129228\n",
      " 10/10... rate=0.00 Hz, eta=0:00:00, total=2:06:02, wall=00:56 UTC\n",
      "  0/10... rate=0 Hz, eta=?, total=0:00:00, wall=00:57 UTCouter fold:4, epoch:0\n",
      "t = 100, loss = 2.42146277\n",
      "t = 200, loss = 1.71867001\n",
      "t = 300, loss = 1.33924687\n",
      "t = 400, loss = 1.19193411\n",
      "t = 500, loss = 1.11659217\n",
      "t = 600, loss = 1.11242020\n",
      "t = 700, loss = 1.05968821\n",
      "average validation loss: 1.02530074\n",
      "  1/10... rate=0.00 Hz, eta=1:53:21, total=0:12:35, wall=01:10 UTCouter fold:4, epoch:1\n",
      "t = 100, loss = 1.05561984\n",
      "t = 200, loss = 1.01083839\n",
      "t = 300, loss = 0.98240125\n",
      "t = 400, loss = 0.99849975\n",
      "t = 500, loss = 0.96065128\n",
      "t = 600, loss = 0.96661824\n",
      "t = 700, loss = 0.99267089\n",
      "average validation loss: 0.97273219\n",
      "  2/10... rate=0.00 Hz, eta=1:40:48, total=0:25:12, wall=01:22 UTCouter fold:4, epoch:2\n",
      "t = 100, loss = 0.93936217\n",
      "t = 200, loss = 0.96657968\n",
      "t = 300, loss = 0.97209775\n",
      "t = 400, loss = 0.98609316\n",
      "t = 500, loss = 0.95763004\n",
      "t = 600, loss = 0.91906780\n",
      "t = 700, loss = 0.96602088\n",
      "average validation loss: 0.97163373\n",
      "  3/10... rate=0.00 Hz, eta=1:28:13, total=0:37:48, wall=01:35 UTCouter fold:4, epoch:3\n",
      "t = 100, loss = 0.96087968\n",
      "t = 200, loss = 0.93622810\n",
      "t = 300, loss = 0.95992321\n",
      "t = 400, loss = 0.95069039\n",
      "t = 500, loss = 0.91604638\n",
      "t = 600, loss = 0.98685670\n",
      "t = 700, loss = 0.94705600\n",
      "average validation loss: 0.97130466\n",
      "  4/10... rate=0.00 Hz, eta=1:15:36, total=0:50:24, wall=01:47 UTCouter fold:4, epoch:4\n",
      "t = 100, loss = 0.95758569\n",
      "t = 200, loss = 0.96495241\n",
      "t = 300, loss = 0.95174968\n",
      "t = 400, loss = 0.99809963\n",
      "t = 500, loss = 0.92730242\n",
      "t = 600, loss = 0.98101997\n",
      "t = 700, loss = 0.96778274\n",
      "average validation loss: 0.96368688\n",
      "  5/10... rate=0.00 Hz, eta=1:03:00, total=1:03:00, wall=02:00 UTCouter fold:4, epoch:5\n",
      "t = 100, loss = 0.90852404\n",
      "t = 200, loss = 0.97151297\n",
      "t = 300, loss = 0.96296978\n",
      "t = 400, loss = 0.99504751\n",
      "t = 500, loss = 1.04426372\n",
      "t = 600, loss = 1.01945925\n",
      "t = 700, loss = 1.02655089\n",
      "average validation loss: 0.98775381\n",
      "  6/10... rate=0.00 Hz, eta=0:50:24, total=1:15:36, wall=02:13 UTCouter fold:4, epoch:6\n",
      "t = 100, loss = 0.95161915\n",
      "t = 200, loss = 1.02700949\n",
      "t = 300, loss = 0.97480845\n",
      "t = 400, loss = 1.00430250\n",
      "t = 500, loss = 0.96381927\n",
      "t = 600, loss = 1.00610495\n",
      "t = 700, loss = 1.00099432\n",
      "average validation loss: 0.99020302\n",
      "  7/10... rate=0.00 Hz, eta=0:37:48, total=1:28:12, wall=02:25 UTCouter fold:4, epoch:7\n",
      "t = 100, loss = 0.99101824\n",
      "t = 200, loss = 1.03180015\n",
      "t = 300, loss = 1.00738609\n",
      "t = 400, loss = 1.03916669\n",
      "t = 500, loss = 1.00257719\n",
      "t = 600, loss = 0.98611605\n",
      "t = 700, loss = 0.98884135\n",
      "average validation loss: 1.00504398\n",
      "  8/10... rate=0.00 Hz, eta=0:25:12, total=1:40:48, wall=02:38 UTCouter fold:4, epoch:8\n",
      "t = 100, loss = 1.08634901\n",
      "t = 200, loss = 1.00054920\n",
      "t = 300, loss = 0.97500402\n",
      "t = 400, loss = 0.97679740\n",
      "t = 500, loss = 0.98815984\n",
      "t = 600, loss = 1.06122482\n",
      "t = 700, loss = 1.01996195\n",
      "average validation loss: 1.00601935\n",
      "  9/10... rate=0.00 Hz, eta=0:12:36, total=1:53:24, wall=02:50 UTCouter fold:4, epoch:9\n",
      "t = 100, loss = 0.94817591\n",
      "t = 200, loss = 0.87107658\n",
      "t = 300, loss = 0.88924980\n",
      "t = 400, loss = 0.91463035\n",
      "t = 500, loss = 0.90386182\n",
      "t = 600, loss = 0.89577562\n",
      "t = 700, loss = 0.89567369\n",
      "average validation loss: 0.92307913\n",
      " 10/10... rate=0.00 Hz, eta=0:00:00, total=2:06:00, wall=03:03 UTC\n"
     ]
    }
   ],
   "source": [
    "batch_size = 1900\n",
    "patience = 2\n",
    "n_epochs = 10\n",
    "dtype=torch.cuda.FloatTensor\n",
    "charset = processor.charset\n",
    "latent_space_size = 16\n",
    "\n",
    "best = torch.load(f'/tf/notebooks/code_for_pub/_logs_as_python_files/vae_training_logs/{str(latent_space_size)}_model_best.pth.tar')\n",
    "#encoder.apply(VAE_model.initialize_weights)\n",
    "#decoder = VAE_model.MolDecoder(i=latent_space_size, o=140, c=len(charset))\n",
    "#decoder.apply(VAE_model.initialize_weights)\n",
    "encoder = VAE_model.MolEncoder(i=140, o=latent_space_size, c=len(charset))\n",
    "encoder.load_state_dict(best['encoder'])\n",
    "\n",
    "decoder = VAE_model.MolDecoder(i=latent_space_size, o=140, c=len(charset))\n",
    "decoder.load_state_dict(best['decoder'])\n",
    "\n",
    "encoder.cuda()\n",
    "decoder.cuda()\n",
    "best_loss = 1e6\n",
    "\n",
    "optimizer = torch.optim.Adam(chain(encoder.parameters(), decoder.parameters()), lr=1e-3)\n",
    "scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, \n",
    "                                                       factor=0.8,\n",
    "                                                       patience=patience,\n",
    "                                                       mode='min', \n",
    "                                                       min_lr=1e-6)\n",
    "\n",
    "kf = KFold(n_splits=5, shuffle=True, random_state=random_state)        \n",
    "for ind,(train_index, test_index) in enumerate(kf.split(df)):\n",
    "    \n",
    "    # make loader train\n",
    "    dt = torch.from_numpy(df[train_index])\n",
    "    train = TensorDataset(dt, torch.zeros(dt.size()[0]))\n",
    "    train_loader = DataLoader(train, batch_size=batch_size, shuffle=True)\n",
    "\n",
    "    # make loader test\n",
    "    dte = torch.from_numpy(df[test_index])\n",
    "    test = TensorDataset(dte, torch.zeros(dte.size()[0]))\n",
    "    val_loader = DataLoader(test, batch_size=batch_size, shuffle=True)\n",
    "    del dt, dte, train, test\n",
    "    \n",
    "    # for other folds we re-use previous best model's weights\n",
    "    if ind != 0:\n",
    "        encoder = VAE_model.MolEncoder(i=140, o=latent_space_size, c=len(charset))\n",
    "        best = torch.load(f'/tf/notebooks/code_for_pub/_logs_as_python_files/vae_training_logs/{str(latent_space_size)}_model_best.pth.tar')\n",
    "        encoder.load_state_dict(best['encoder'])\n",
    "        decoder = VAE_model.MolDecoder(i=latent_space_size, o=140, c=len(charset))\n",
    "        decoder.load_state_dict(best['decoder'])\n",
    "        encoder.cuda()\n",
    "        decoder.cuda()\n",
    "\n",
    "        optimizer = torch.optim.Adam(chain(encoder.parameters(), decoder.parameters()), lr=1e-3)\n",
    "        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, \n",
    "                                                               factor=0.2,\n",
    "                                                               patience=patience,\n",
    "                                                               mode='min', \n",
    "                                                               min_lr=1e-6)\n",
    "    \n",
    "    \n",
    "    \n",
    "    # train for n epochs\n",
    "    for epoch in ProgIter(range(n_epochs), verbose=1, total=n_epochs):\n",
    "        print(f'outer fold:{ind}, epoch:{epoch}')\n",
    "        VAE_model.train_model(train_loader, encoder, decoder, optimizer, dtype)\n",
    "        avg_val_loss = VAE_model.validate_model(val_loader, encoder, decoder, dtype)\n",
    "\n",
    "        scheduler.step(avg_val_loss)\n",
    "\n",
    "        is_best = avg_val_loss < best_loss\n",
    "\n",
    "        if is_best:\n",
    "            best_loss = avg_val_loss\n",
    "        VAE_model.save_checkpoint({\n",
    "            'epoch': epoch,\n",
    "            'encoder': encoder.state_dict(),\n",
    "            'decoder': decoder.state_dict(),\n",
    "            'charset': charset,\n",
    "            'avg_val_loss': avg_val_loss,\n",
    "            'optimizer': optimizer.state_dict(),\n",
    "        }, is_best, size = latent_space_size)\n",
    "\n",
    "    del train_loader, encoder, decoder, scheduler, optimizer \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:root] *",
   "language": "python",
   "name": "conda-root-py"
  },
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}