ex_cosmology/01_train_models.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "2rg2KQ-iieHN"
   },
   "source": [
    "## Data Preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "Akr87A3ojAkR",
    "outputId": "5ba4d687-165b-4661-ac05-ee78bcb356c1"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Populating the interactive namespace from numpy and matplotlib\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import os\n",
    "import torch\n",
    "from skimage import io, transform\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torchvision import transforms, utils\n",
    "from astropy.io import fits\n",
    "import torchvision\n",
    "import torch.nn as nn\n",
    "from torchvision import datasets, models, transforms\n",
    "import time\n",
    "import os\n",
    "import copy\n",
    "from copy import deepcopy\n",
    "from sklearn.metrics import r2_score\n",
    "from data import MassMapsDatasetFilteredS8 as MassMapsDataset\n",
    "from data import RandomToTensor\n",
    "import data\n",
    "import torchvision.transforms.functional as F\n",
    "import random\n",
    "from os.path import join as oj\n",
    "import torch.optim as optim\n",
    "from torch.optim import lr_scheduler\n",
    "plt.style.use('dark_background')\n",
    "out_dir = '/scratch/users/vision/data/cosmo/'\n",
    "\n",
    "# Ignore warnings\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "%pylab inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# download / extract data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "gs://massivenu/cosmological_parameters.txt\n",
      "gs://massivenu/resnet18\n",
      "gs://massivenu/resnet18_state_dict\n",
      "gs://massivenu/models_feb20/\n",
      "gs://massivenu/z1_256/\n",
      "gs://massivenu/z1_256b/\n",
      "gs://massivenu/z1_256b_test/\n"
     ]
    }
   ],
   "source": [
    "!gsutil ls gs://massivenu\n",
    "folders = ['z1_256b', 'z1_256b_test']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cosmological_parameters.txt  resnet18_state_dict  z1_256   z1_256b_test\n",
      "resnet18\t\t     results\t\t  z1_256b\n"
     ]
    }
   ],
   "source": [
    "!ls {out_dir}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "id": "snWLaALChtBc",
    "outputId": "86c6f27c-6acb-45f9-f0d3-225a39fbd348"
   },
   "outputs": [],
   "source": [
    "# First step, download data, it takes a few minutes....\n",
    "# ! gsutil -m cp -r gs://massivenu/cosmological_parameters.txt {out_dir}\n",
    "# for folder in folders:\n",
    "#     ! gsutil -m cp -r {'gs://massivenu/' + folder} {out_dir}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "# Second step, extract data, that also takes a few minutes\n",
    "# ! cd z1_256b ; for m in model*.tar.gz; do tar -xzf $m ; rm $m ; done;\n",
    "cd '/scratch/users/vision/data/cosmo/z1_256b_test'\n",
    "# for m in model*.tar.gz\n",
    "# do\n",
    "#     tar -xzf $m\n",
    "#     rm $m\n",
    "# done"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# make dset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Cy7CrPu1w-yM"
   },
   "outputs": [],
   "source": [
    "data_transform = transforms.Compose([RandomToTensor()])\n",
    "param_file = oj(out_dir, 'cosmological_parameters.txt')\n",
    "mnu_dataset = MassMapsDataset(param_file,  \n",
    "                              root_dir=oj(out_dir, 'z1_256b'),\n",
    "                              transform=data_transform)\n",
    "mnu_dataset_test = MassMapsDataset(param_file,\n",
    "                              root_dir=oj(out_dir, 'z1_256b_test'),\n",
    "                              transform=data_transform)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**look at some ims**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "id": "BptxJT9loPmo",
    "outputId": "a4f30b94-6a98-4099-910a-0fc150e83d29"
   },
   "outputs": [],
   "source": [
    "figure(figsize=(15,15))\n",
    "mes = []\n",
    "s8 = []\n",
    "om = []\n",
    "for i in range(len(mnu_dataset)):\n",
    "    sample = mnu_dataset[i]\n",
    "    #print(i, sample['image'].shape, sample['params'].shape)\n",
    "\n",
    "    subplot(5, 5, i + 1)\n",
    "    tight_layout()\n",
    "    axis('off')\n",
    "    imshow(squeeze(sample['image']), cmap='magma',vmin=-0.1,vmax=0.15)\n",
    "    title(r\"$M_\\nu=%0.2f ; \\Omega_m$=%0.2f; $\\sigma_8=%0.2f$\"%(sample['params'][0],sample['params'][1],sample['params'][2] ) )\n",
    "    s8.append(sample['params'][2].numpy())\n",
    "    om.append(sample['params'][1].numpy())\n",
    "    mes.append(np.std(sample['image'].numpy()))\n",
    "    if i == 24:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "PpcDjnVSqQDH",
    "outputId": "137dae4b-023b-442c-91ff-dde0268d9252"
   },
   "outputs": [],
   "source": [
    "params = mnu_dataset.params\n",
    "R, C = 1, 5\n",
    "plt.figure(dpi=200, figsize=(7, 2))\n",
    "for i in range(5):\n",
    "    plt.subplot(R, C, i + 1)\n",
    "    plt.hist(params[:, i], bins=20)\n",
    "    plt.title(data.classes[i])\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Nfvz34PiqfCt"
   },
   "source": [
    "## Train a model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "FhEY-N9SvYT9",
    "outputId": "c4276266-5710-40d5-d3ae-0c3e61f4ac04"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda:0\n"
     ]
    }
   ],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "fhD9pDs1M3jy"
   },
   "outputs": [],
   "source": [
    "# Modifying the model to predict the three cosmological parameters from single channel images\n",
    "model_ft = models.resnet18(pretrained=False)\n",
    "model_ft.conv1 = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
    "num_ftrs = model_ft.fc.in_features\n",
    "model_ft.fc = nn.Linear(num_ftrs, 3)\n",
    "model_ft = model_ft.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Fk_0M-AurPCD"
   },
   "outputs": [],
   "source": [
    "# data_transform = transforms.Compose([\n",
    "#         ToTensor()\n",
    "#     ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "s1tPV8_RrPG0"
   },
   "outputs": [],
   "source": [
    "# mnu_dataset = MassMapsDataset('cosmological_parameters.txt',\n",
    "#                               'z1_256b', \n",
    "#                               transform=data_transform)\n",
    "\n",
    "dataloader = torch.utils.data.DataLoader(mnu_dataset, batch_size=128, \n",
    "                                         shuffle=True, num_workers=4)\n",
    "dataloader_test = torch.utils.data.DataLoader(mnu_dataset_test, batch_size=128, \n",
    "                                         shuffle=False, num_workers=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "LilJJ0Nsuk3J"
   },
   "outputs": [],
   "source": [
    "def train_model(model, dataloader_train, dataloader_test, criterion, optimizer, scheduler, num_epochs=25):\n",
    "    since = time.time()\n",
    "\n",
    "    best_model_wts = deepcopy(model.state_dict())\n",
    "    best_loss = 1e10\n",
    "    train_losses = []\n",
    "    test_losses = []\n",
    "    learning_rates= []\n",
    "    for epoch in range(num_epochs):\n",
    "        print('Epoch {}/{}'.format(epoch, num_epochs - 1))\n",
    "        print('-' * 10)\n",
    "\n",
    "        # Each epoch has a training and validation phase\n",
    "        for phase in ['train', 'test']:\n",
    "            if phase == 'train':\n",
    "                model.train()  # Set model to training mode\n",
    "                dataloader = dataloader_train\n",
    "            else:\n",
    "                model.eval()   # Set model to evaluate mode\n",
    "                dataloader = dataloader_test\n",
    "\n",
    "            running_loss = 0.0\n",
    "            \n",
    "            # Iterate over data.\n",
    "            for data in dataloader:\n",
    "                inputs, params = data['image'], data['params']\n",
    "                inputs = inputs.to(device)\n",
    "                params = params.to(device)\n",
    "\n",
    "                # zero the parameter gradients\n",
    "                optimizer.zero_grad()\n",
    "\n",
    "                # forward\n",
    "                # track history if only in train\n",
    "                with torch.set_grad_enabled(phase == 'train'):\n",
    "                    outputs = model(inputs)\n",
    "                    loss = criterion(outputs, params)\n",
    "                    \n",
    "                    # backward + optimize only if in training phase\n",
    "                    if phase == 'train':\n",
    "                        loss.backward()\n",
    "                        optimizer.step()\n",
    "\n",
    "                # statistics\n",
    "                running_loss += loss.item() * inputs.size(0)\n",
    "                if phase == 'train':\n",
    "                    train_losses.append(loss.item())\n",
    "                else:\n",
    "                    test_losses.append(loss.item())\n",
    "            if phase == 'train':\n",
    "                scheduler.step()\n",
    "                learning_rates.append(scheduler.get_lr())\n",
    "\n",
    "            epoch_loss = running_loss / len(mnu_dataset)\n",
    "            \n",
    "            print('{} Loss: {:.4f} '.format(\n",
    "                phase, epoch_loss))\n",
    "            \n",
    "            if epoch_loss <= best_loss:\n",
    "                torch.save(deepcopy(model_ft.state_dict()), oj(out_dir, f'resnet18_full_state_dict_{epoch}.pkl'))\n",
    "            best_loss = min(epoch_loss, best_loss)\n",
    "            \n",
    "            \n",
    "            # save\n",
    "            # Saving fairly well trained model\n",
    "            \n",
    "            # # deep copy the model\n",
    "            # if phase == 'val' and epoch_loss < best_loss:\n",
    "            #     best_loss = epoch_loss\n",
    "            #     best_model_wts = copy.deepcopy(model.state_dict())\n",
    "\n",
    "    time_elapsed = time.time() - since\n",
    "    print('Training complete in {:.0f}m {:.0f}s'.format(\n",
    "        time_elapsed // 60, time_elapsed % 60))\n",
    "    print('Best val Loss: {:4f}'.format(best_loss))\n",
    "\n",
    "    # load best model weights\n",
    "    # model.load_state_dict(best_model_wts)\n",
    "    return model, train_losses, test_losses, learning_rates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Wah3aOO9w3Ib"
   },
   "outputs": [],
   "source": [
    "criterion = torch.nn.L1Loss()\n",
    "\n",
    "# Observe that all parameters are being optimized\n",
    "optimizer_ft = optim.Adam(model_ft.parameters(), lr=0.001)\n",
    "\n",
    "# Decay LR by a factor of 0.1 every 7 epochs\n",
    "exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=10, gamma=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "id": "cb-Kx4jW3nmD",
    "outputId": "e105b486-dd56-4755-878a-5c8f5f7d3cf1"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0/24\n",
      "----------\n",
      "train Loss: 0.0739 \n"
     ]
    },
    {
     "ename": "FileNotFoundError",
     "evalue": "Caught FileNotFoundError in DataLoader worker process 1.\nOriginal Traceback (most recent call last):\n  File \"/accounts/projects/vision/.local/lib/python3.7/site-packages/torch/utils/data/_utils/worker.py\", line 178, in _worker_loop\n    data = fetcher.fetch(index)\n  File \"/accounts/projects/vision/.local/lib/python3.7/site-packages/torch/utils/data/_utils/fetch.py\", line 44, in fetch\n    data = [self.dataset[idx] for idx in possibly_batched_index]\n  File \"/accounts/projects/vision/.local/lib/python3.7/site-packages/torch/utils/data/_utils/fetch.py\", line 44, in <listcomp>\n    data = [self.dataset[idx] for idx in possibly_batched_index]\n  File \"/accounts/projects/vision/chandan/transformation-interpretability/ex_cosmology/data.py\", line 94, in __getitem__\n    image = fits.getdata(img_name)\n  File \"/accounts/projects/vision/.local/lib/python3.7/site-packages/astropy/io/fits/convenience.py\", line 189, in getdata\n    hdulist, extidx = _getext(filename, mode, *args, **kwargs)\n  File \"/accounts/projects/vision/.local/lib/python3.7/site-packages/astropy/io/fits/convenience.py\", line 1029, in _getext\n    hdulist = fitsopen(filename, mode=mode, **kwargs)\n  File \"/accounts/projects/vision/.local/lib/python3.7/site-packages/astropy/io/fits/hdu/hdulist.py\", line 151, in fitsopen\n    lazy_load_hdus, **kwargs)\n  File \"/accounts/projects/vision/.local/lib/python3.7/site-packages/astropy/io/fits/hdu/hdulist.py\", line 390, in fromfile\n    lazy_load_hdus=lazy_load_hdus, **kwargs)\n  File \"/accounts/projects/vision/.local/lib/python3.7/site-packages/astropy/io/fits/hdu/hdulist.py\", line 1039, in _readfrom\n    fileobj = _File(fileobj, mode=mode, memmap=memmap, cache=cache)\n  File \"/accounts/projects/vision/.local/lib/python3.7/site-packages/astropy/utils/decorators.py\", line 521, in wrapper\n    return function(*args, **kwargs)\n  File \"/accounts/projects/vision/.local/lib/python3.7/site-packages/astropy/io/fits/file.py\", line 178, in __init__\n    self._open_filename(fileobj, mode, overwrite)\n  File \"/accounts/projects/vision/.local/lib/python3.7/site-packages/astropy/io/fits/file.py\", line 555, in _open_filename\n    self._file = fileobj_open(self.name, IO_FITS_MODES[mode])\n  File \"/accounts/projects/vision/.local/lib/python3.7/site-packages/astropy/io/fits/util.py\", line 397, in fileobj_open\n    return open(filename, mode, buffering=0)\nFileNotFoundError: [Errno 2] No such file or directory: '/scratch/users/vision/data/cosmo/z1_256b_test/model003/WLconv_z1.00_0100r.fits'\n",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mFileNotFoundError\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-11-7255c140ddd0>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m model_ft, loss_trace, lr_trace = train_model(model_ft, dataloader, dataloader_test,\n\u001b[0;32m----> 2\u001b[0;31m                                              criterion, optimizer_ft, exp_lr_scheduler, num_epochs=25)\n\u001b[0m",
      "\u001b[0;32m<ipython-input-9-6e6f32899273>\u001b[0m in \u001b[0;36mtrain_model\u001b[0;34m(model, dataloader_train, dataloader_test, criterion, optimizer, scheduler, num_epochs)\u001b[0m\n\u001b[1;32m     23\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     24\u001b[0m             \u001b[0;31m# Iterate over data.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 25\u001b[0;31m             \u001b[0;32mfor\u001b[0m \u001b[0mdata\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mdataloader\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     26\u001b[0m                 \u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparams\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdata\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'image'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdata\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'params'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     27\u001b[0m                 \u001b[0minputs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.7/site-packages/torch/utils/data/dataloader.py\u001b[0m in \u001b[0;36m__next__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    817\u001b[0m             \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    818\u001b[0m                 \u001b[0;32mdel\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_task_info\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0midx\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 819\u001b[0;31m                 \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_process_data\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdata\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    820\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    821\u001b[0m     \u001b[0mnext\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m__next__\u001b[0m  \u001b[0;31m# Python 2 compatibility\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.7/site-packages/torch/utils/data/dataloader.py\u001b[0m in \u001b[0;36m_process_data\u001b[0;34m(self, data)\u001b[0m\n\u001b[1;32m    844\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_try_put_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    845\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mExceptionWrapper\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 846\u001b[0;31m             \u001b[0mdata\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreraise\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    847\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mdata\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    848\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.7/site-packages/torch/_utils.py\u001b[0m in \u001b[0;36mreraise\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    383\u001b[0m             \u001b[0;31m# (https://bugs.python.org/issue2651), so we work around it.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    384\u001b[0m             \u001b[0mmsg\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mKeyErrorMessage\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmsg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 385\u001b[0;31m         \u001b[0;32mraise\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexc_type\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmsg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mFileNotFoundError\u001b[0m: Caught FileNotFoundError in DataLoader worker process 1.\nOriginal Traceback (most recent call last):\n  File \"/accounts/projects/vision/.local/lib/python3.7/site-packages/torch/utils/data/_utils/worker.py\", line 178, in _worker_loop\n    data = fetcher.fetch(index)\n  File \"/accounts/projects/vision/.local/lib/python3.7/site-packages/torch/utils/data/_utils/fetch.py\", line 44, in fetch\n    data = [self.dataset[idx] for idx in possibly_batched_index]\n  File \"/accounts/projects/vision/.local/lib/python3.7/site-packages/torch/utils/data/_utils/fetch.py\", line 44, in <listcomp>\n    data = [self.dataset[idx] for idx in possibly_batched_index]\n  File \"/accounts/projects/vision/chandan/transformation-interpretability/ex_cosmology/data.py\", line 94, in __getitem__\n    image = fits.getdata(img_name)\n  File \"/accounts/projects/vision/.local/lib/python3.7/site-packages/astropy/io/fits/convenience.py\", line 189, in getdata\n    hdulist, extidx = _getext(filename, mode, *args, **kwargs)\n  File \"/accounts/projects/vision/.local/lib/python3.7/site-packages/astropy/io/fits/convenience.py\", line 1029, in _getext\n    hdulist = fitsopen(filename, mode=mode, **kwargs)\n  File \"/accounts/projects/vision/.local/lib/python3.7/site-packages/astropy/io/fits/hdu/hdulist.py\", line 151, in fitsopen\n    lazy_load_hdus, **kwargs)\n  File \"/accounts/projects/vision/.local/lib/python3.7/site-packages/astropy/io/fits/hdu/hdulist.py\", line 390, in fromfile\n    lazy_load_hdus=lazy_load_hdus, **kwargs)\n  File \"/accounts/projects/vision/.local/lib/python3.7/site-packages/astropy/io/fits/hdu/hdulist.py\", line 1039, in _readfrom\n    fileobj = _File(fileobj, mode=mode, memmap=memmap, cache=cache)\n  File \"/accounts/projects/vision/.local/lib/python3.7/site-packages/astropy/utils/decorators.py\", line 521, in wrapper\n    return function(*args, **kwargs)\n  File \"/accounts/projects/vision/.local/lib/python3.7/site-packages/astropy/io/fits/file.py\", line 178, in __init__\n    self._open_filename(fileobj, mode, overwrite)\n  File \"/accounts/projects/vision/.local/lib/python3.7/site-packages/astropy/io/fits/file.py\", line 555, in _open_filename\n    self._file = fileobj_open(self.name, IO_FITS_MODES[mode])\n  File \"/accounts/projects/vision/.local/lib/python3.7/site-packages/astropy/io/fits/util.py\", line 397, in fileobj_open\n    return open(filename, mode, buffering=0)\nFileNotFoundError: [Errno 2] No such file or directory: '/scratch/users/vision/data/cosmo/z1_256b_test/model003/WLconv_z1.00_0100r.fits'\n"
     ]
    }
   ],
   "source": [
    "model_ft, loss_trace, lr_trace = train_model(model_ft, dataloader, dataloader_test,\n",
    "                                             criterion, optimizer_ft, exp_lr_scheduler, num_epochs=25)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "g4TdOt01dWve"
   },
   "outputs": [],
   "source": [
    "# # Observe that all parameters are being optimized\n",
    "# optimizer_ft = optim.Adam(model_ft.parameters(), lr=0.0001)\n",
    "\n",
    "# # Decay LR by a factor of 0.1 every 7 epochs\n",
    "# exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=10, gamma=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "id": "kGkm7zVLdij4",
    "outputId": "fac7615f-73e8-4ede-8f0c-b1d87b2734e4"
   },
   "outputs": [],
   "source": [
    "# model_ft, loss_trace2, lr_trace2 = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler, num_epochs=25)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# analyze model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "7ST4boEbfTe_"
   },
   "outputs": [],
   "source": [
    "plot(loss_trace2[100:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "1uVup359DgKs"
   },
   "outputs": [],
   "source": [
    "plot(lr_trace2)\n",
    "yscale('log')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "NBueEXv6DgPU"
   },
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    for data in dataloader:\n",
    "        inputs, params = data['image'], data['params']\n",
    "        inputs = inputs.to(device)\n",
    "        params = params.to(device)\n",
    "        outputs = model_ft(inputs)\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 279
    },
    "colab_type": "code",
    "id": "6z9JHvbkDgI0",
    "outputId": "0a967858-e263-4767-da36-09c8a071e0f1"
   },
   "outputs": [],
   "source": [
    "scatter(params.cpu()[:,1], outputs.cpu()[:,1])\n",
    "xlabel('true param')\n",
    "ylabel('predicted param')\n",
    "#ylim(0.265,0.35)\n",
    "#xlim(0.265,0.35)\n",
    "gca().set_aspect('equal')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "D3Svf_TADgHO"
   },
   "outputs": [],
   "source": [
    "coefficient_of_dermination = r2_score(params.cpu()[:,1], outputs.cpu()[:,1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "Yx77qU6WDgED",
    "outputId": "20d9c6ec-5087-431c-a04e-2c0050710f16"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.7257992068325176"
      ]
     },
     "execution_count": 71,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "coefficient_of_dermination"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "F2FAOFti3Qte"
   },
   "outputs": [],
   "source": [
    "# Saving fairly well trained model\n",
    "torch.save(model_ft.state_dict(), 'resnet18_state_dict')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "oT95WipVXfAs"
   },
   "outputs": [],
   "source": [
    "torch.save(model_ft, 'resnet18')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "KlwSJgjiPnLc"
   },
   "outputs": [],
   "source": [
    "from google.colab import auth\n",
    "auth.authenticate_user()"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "DeepMassACDb.ipynb",
   "provenance": []
  },
  "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}