final commit to hf

README.md

@@ -0,0 +1,5 @@
+# How to run locally
+
+- [ ] ```git clone https://github.com/majauhar/UpsideDownDetector.git```
+- [ ] ```pip install -r requirements.txt```
+- [ ] ```python main.py --epochs=<Integer> --pretrained=[True/False]```

Viz.ipynb

@@ -0,0 +1,51 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "ModuleNotFoundError",
+     "evalue": "No module named 'torch'",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
+      "\u001b[0;32m<ipython-input-5-eb42ca6e4af3>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0;32mimport\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
+      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'torch'"
+     ]
+    }
+   ],
+   "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.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

main.py

@@ -0,0 +1,100 @@
+import torch
+import torchvision
+from torchvision import datasets, transforms
+from torch.utils.data import Dataset, DataLoader
+from torchvision import models
+import numpy as np
+import matplotlib.pyplot as plt
+from PIL import Image
+import os
+import os.path
+import pickle
+from typing import Any, Callable, Optional, Tuple
+import torchvision.transforms.functional as TF
+
+from tools.dataset import UpsideDownDataset
+from train import train
+from test import test
+from observations import observations
+
+import torchvision.models as models
+import torch.optim as optim
+
+import argparse
+
+# setting  up the device
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+def arg_parser():
+    parser = argparse.ArgumentParser(description="Upside Down Image Detector")
+    parser.add_argument('--epochs', type=int, default=10)
+    parser.add_argument('--pretrained', type=bool, default=True)
+
+    return parser.parse_args()
+
+def main(args):
+    # transformations
+    train_transform = transforms.Compose([
+        transforms.ToTensor()
+    ])
+
+    test_transform = transforms.Compose([
+        transforms.ToTensor()
+    ])
+
+    # downloading CIFAR10 dataset
+    trainset = UpsideDownDataset(root='./data', train=True, download=True, transform=train_transform)
+    testset = UpsideDownDataset(root='./data', train=False, download=True, transform=test_transform)
+
+    trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, drop_last=True)
+    testloader = torch.utils.data.DataLoader(testset, batch_size=128)
+
+
+    model = models.resnet18(pretrained=args.pretrained)
+    model.fc = torch.nn.Linear(512, 2)
+
+    optimizer = optim.SGD(model.parameters(), lr=0.01)
+    # criterion = torch.nn.BCEWithLogitsLoss()
+    criterion = torch.nn.CrossEntropyLoss()
+
+    model, criterion = model.to(DEVICE), criterion.to(DEVICE)
+
+    print("training started:")
+
+    loss_values = []
+    accuracy_values = []
+    for epoch in range(args.epochs):
+        print("Epoch {}".format(epoch+1))
+        epoch_loss = train(model, trainloader, optimizer, criterion, DEVICE)
+        loss_values.append(epoch_loss)
+        accuracy = test(model, testloader, criterion, DEVICE)
+        accuracy_values.append(accuracy)
+
+    plt.plot(np.array(loss_values), 'r')
+    plt.xlabel('epochs')
+    plt.ylabel('losses')
+    plt.savefig('./viz/loss_curve.png')
+
+    fig, ax = plt.subplots()
+    ax.plot(np.array(loss_values), color='red', label='losses')
+    ax.tick_params(axis='y', labelcolor='red')
+
+    ax2 = ax.twinx()
+
+    ax2.plot(np.array(accuracy_values), color='green', label='accuracy')
+    ax2.tick_params(axis='y', labelcolor='green')
+
+    # plt.xlabel('epochs')
+    plt.legend()
+    plt.savefig('./viz/final_curve.png')
+
+    print("accuracy: {}".format(accuracy))
+
+    observations(model, testloader)
+
+
+
+
+if __name__ =='__main__':
+    args = arg_parser()
+    main(args)

observations.py

@@ -0,0 +1,54 @@
+import torch
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+def pred_label(model, img):
+    model = model.to('cpu')
+    img = img.unsqueeze(0)
+    logits = model(img)
+    pred_probab = torch.nn.Softmax(dim=1)(logits)
+    y_pred = pred_probab.argmax(1)
+
+    return y_pred
+
+def save_image(img, title, count):
+    fig, ax = plt.subplots()
+    imgplot = ax.imshow(img, interpolation='bicubic')
+    ax.spines['top'].set_visible(False)
+    ax.spines['left'].set_visible(False)
+    ax.spines['bottom'].set_visible(False)
+    ax.spines['right'].set_visible(False)
+    ax.set_xticks([])
+    ax.set_yticks([])
+    # imgplot = plt.imshow(img, interpolation='bicubic')
+    plt.title(title)
+    plt.savefig('./viz/img' + str(count))
+
+
+def observations(model, testloader):
+    for imgs, labels in testloader:
+        images = [imgs[0].permute(1, 2, 0),
+        imgs[1].permute(1, 2, 0),
+        imgs[2].permute(1, 2, 0),
+        imgs[3].permute(1, 2, 0),
+        imgs[4].permute(1, 2, 0)]
+
+        pred_label1 = pred_label(model, imgs[0]).item()
+        pred_label2 = pred_label(model, imgs[1]).item()
+        pred_label3 = pred_label(model, imgs[2]).item()
+        pred_label4 = pred_label(model, imgs[3]).item()
+        pred_label5 = pred_label(model, imgs[4]).item()
+
+        titles = ["Pred: {}, Actual: {}".format(pred_label1, labels[0]), 
+        "Pred: {}, Actual: {}".format(pred_label2, labels[1]),
+        "Pred: {}, Actual: {}".format(pred_label3, labels[2]),
+        "Pred: {}, Actual: {}".format(pred_label4, labels[3]),
+        "Pred: {}, Actual: {}".format(pred_label5, labels[4])]
+
+        count = 1
+        for image, title in zip(images, titles):
+            save_image(image, title, count)
+            count += 1
+
+        break

requirements.txt

@@ -0,0 +1,285 @@
+absl-py==0.10.0
+alabaster==0.7.12
+allennlp==0.8.4
+anaconda-client==1.7.2
+anaconda-navigator==1.9.7
+anaconda-project==0.8.3
+asn1crypto==0.24.0
+astroid==2.2.5
+astropy==3.2.1
+astunparse==1.6.3
+atomicwrites==1.3.0
+attrs==20.1.0
+awscli==1.18.132
+Babel==2.8.0
+backcall==0.1.0
+backports.functools-lru-cache==1.5
+backports.os==0.1.1
+backports.shutil-get-terminal-size==1.0.0
+backports.tempfile==1.0
+backports.weakref==1.0.post1
+beautifulsoup4==4.7.1
+bitarray==0.9.3
+bkcharts==0.2
+bleach==3.1.0
+blis==0.2.4
+bokeh==1.2.0
+boto==2.49.0
+boto3==1.14.55
+botocore==1.17.55
+Bottleneck==1.2.1
+cachetools==4.1.1
+certifi==2021.5.30
+cffi==1.12.3
+chardet==3.0.4
+click==7.1.2
+cloudpickle==1.2.1
+clyent==1.2.2
+colorama==0.4.3
+conda==4.7.10
+conda-build==3.18.8
+conda-package-handling==1.3.11
+conda-verify==3.4.2
+conllu==0.11
+contextlib2==0.5.5
+cryptography==2.7
+cycler==0.10.0
+Cython==0.29.12
+cytoolz==0.10.0
+dask==2.1.0
+decorator==4.4.0
+defusedxml==0.6.0
+detectron2==0.2.1+cu101
+distributed==2.1.0
+docutils==0.15.2
+editdistance==0.5.3
+entrypoints==0.3
+et-xmlfile==1.0.1
+fastapi==0.63.0
+fastcache==1.1.0
+filelock==3.0.12
+flaky==3.7.0
+Flask==1.1.2
+Flask-Cors==3.0.9
+ftfy==5.8
+future==0.17.1
+fvcore==0.1.1.post20200716
+gast==0.3.3
+gdown==3.13.0
+gevent==20.6.2
+glob2==0.7
+gmpy2==2.0.8
+google-auth==1.20.1
+google-auth-oauthlib==0.4.1
+google-pasta==0.2.0
+greenlet==0.4.16
+grpcio==1.31.0
+gunicorn==20.0.4
+h11==0.12.0
+h5py==2.10.0
+heapdict==1.0.0
+html5lib==1.0.1
+idna==2.10
+imageio==2.5.0
+imagesize==1.2.0
+importlib-metadata==1.7.0
+iniconfig==1.0.1
+ipykernel==5.1.1
+ipython==7.6.1
+ipython-genutils==0.2.0
+ipywidgets==7.5.0
+isort==4.3.21
+itsdangerous==1.1.0
+jdcal==1.4.1
+jedi==0.13.3
+jeepney==0.4
+Jinja2==2.11.2
+jmespath==0.10.0
+joblib==0.16.0
+json5==0.8.4
+jsonnet==0.16.0
+jsonpickle==1.4.1
+jsonschema==3.0.1
+jupyter==1.0.0
+jupyter-client==5.3.1
+jupyter-console==6.0.0
+jupyter-core==4.5.0
+jupyterlab==1.0.2
+jupyterlab-server==1.0.0
+Keras-Preprocessing==1.1.2
+keyring==18.0.0
+kiwisolver==1.2.0
+lazy-object-proxy==1.4.1
+libarchive-c==2.8
+lief==0.9.0
+llvmlite==0.29.0
+locket==0.2.0
+lxml==4.3.4
+Markdown==3.2.2
+MarkupSafe==1.1.1
+matplotlib==3.3.1
+mccabe==0.6.1
+mistune==0.8.4
+mkl-fft==1.0.12
+mkl-random==1.0.2
+mkl-service==2.0.2
+mock==3.0.5
+more-itertools==8.5.0
+mpmath==1.1.0
+msgpack==0.6.1
+multipledispatch==0.6.0
+navigator-updater==0.2.1
+nbconvert==5.5.0
+nbformat==4.4.0
+networkx==2.3
+nltk==3.5
+nose==1.3.7
+notebook==6.0.0
+numba==0.44.1
+numexpr==2.6.9
+numpy==1.18.5
+numpydoc==1.1.0
+oauthlib==3.1.0
+olefile==0.46
+openpyxl==2.6.2
+opt-einsum==3.3.0
+overrides==3.1.0
+packaging==20.4
+pandas==0.24.2
+pandocfilters==1.4.2
+parsimonious==0.8.1
+parso==0.5.0
+partd==1.0.0
+path.py==12.0.1
+pathlib2==2.3.4
+patsy==0.5.1
+pep8==1.7.1
+pexpect==4.7.0
+pickleshare==0.7.5
+Pillow==7.2.0
+pkginfo==1.5.0.1
+plac==0.9.6
+pluggy==0.13.1
+ply==3.11
+portalocker==2.0.0
+prometheus-client==0.7.1
+prompt-toolkit==2.0.9
+protobuf==3.13.0
+psutil==5.6.3
+ptyprocess==0.6.0
+py==1.9.0
+pyasn1==0.4.8
+pyasn1-modules==0.2.8
+pycocotools==2.0.1
+pycodestyle==2.5.0
+pycosat==0.6.3
+pycparser==2.19
+pycrypto==2.6.1
+pycurl==7.43.0.3
+pydantic==1.7.3
+pydot==1.4.1
+pyflakes==2.1.1
+Pygments==2.6.1
+pylint==2.3.1
+pyodbc==4.0.26
+pyOpenSSL==19.0.0
+pyparsing==2.4.7
+pyrsistent==0.14.11
+PySocks==1.7.0
+pytesseract==0.3.8
+pytest==6.0.1
+pytest-arraydiff==0.3
+pytest-astropy==0.5.0
+pytest-doctestplus==0.3.0
+pytest-openfiles==0.3.2
+pytest-remotedata==0.3.1
+python-dateutil==2.8.1
+python-Levenshtein==0.12.0
+pytorch-pretrained-bert==0.6.2
+pytorch3d==0.3.0
+pytz==2020.1
+PyWavelets==1.0.3
+PyYAML==5.3.1
+pyzmq==18.0.0
+QtAwesome==0.5.7
+qtconsole==4.5.1
+QtPy==1.8.0
+regex==2020.7.14
+requests==2.24.0
+requests-oauthlib==1.3.0
+responses==0.12.0
+rope==0.14.0
+rsa==4.5
+ruamel-yaml==0.15.46
+s3transfer==0.3.3
+sacremoses==0.0.43
+scikit-image==0.15.0
+scikit-learn==0.20.0
+scipy==1.4.1
+seaborn==0.9.0
+SecretStorage==3.1.1
+Send2Trash==1.5.0
+sentencepiece==0.1.91
+simplegeneric==0.8.1
+singledispatch==3.4.0.3
+six==1.15.0
+snowballstemmer==2.0.0
+sortedcollections==1.1.2
+sortedcontainers==2.1.0
+soupsieve==1.8
+spacy==2.1.9
+Sphinx==3.2.1
+sphinxcontrib-applehelp==1.0.2
+sphinxcontrib-devhelp==1.0.2
+sphinxcontrib-htmlhelp==1.0.3
+sphinxcontrib-jsmath==1.0.1
+sphinxcontrib-qthelp==1.0.3
+sphinxcontrib-serializinghtml==1.1.4
+sphinxcontrib-websupport==1.1.2
+spyder==3.3.6
+spyder-kernels==0.5.1
+SQLAlchemy==1.3.5
+sqlparse==0.3.1
+srsly==1.0.2
+starlette==0.13.6
+statsmodels==0.10.0
+sympy==1.4
+tables==3.5.2
+tabulate==0.8.7
+tblib==1.4.0
+tensorboard==2.3.0
+tensorboard-plugin-wit==1.7.0
+tensorboardX==2.1
+tensorflow-estimator==2.3.0
+tensorflow-gpu==2.3.0
+termcolor==1.1.0
+terminado==0.8.2
+tesseract-ocr==0.0.1
+testpath==0.4.2
+thinc==7.0.8
+toolz==0.10.0
+tornado==6.0.3
+tqdm==4.48.2
+traitlets==4.3.2
+transformers==2.2.2
+typing-extensions==3.7.4.3
+unicodecsv==0.14.1
+Unidecode==1.1.1
+urllib3==1.25.10
+uvicorn==0.13.3
+wasabi==0.8.0
+wcwidth==0.2.5
+webencodings==0.5.1
+Werkzeug==1.0.1
+widgetsnbextension==3.5.0
+word2number==1.1
+wrapt==1.12.1
+wurlitzer==1.0.2
+xlrd==1.2.0
+XlsxWriter==1.1.8
+xlwt==1.3.0
+yacs==0.1.8
+zict==1.0.0
+zipp==3.1.0
+zope.event==4.4
+zope.interface==5.1.0

test.py

@@ -0,0 +1,18 @@
+import torch
+
+def test(model, testloader, criterion, DEVICE):
+    model.eval()
+    test_loss, correct = 0.0, 0
+    
+    with torch.no_grad():
+        for imgs, targets in testloader:
+            imgs, targets = imgs.to(DEVICE), targets.to(DEVICE)
+            pred = model(imgs)
+            loss = criterion(pred, targets)
+            test_loss += loss.item()
+            correct += (pred.argmax(1) == targets).type(torch.float).sum().item()
+    
+    # test_loss = test_loss / len(testloader)
+    accuracy = correct / len(testloader.dataset) * 100
+
+    return accuracy

tools/__init__.py

@@ -0,0 +1 @@
+from . import dataset

tools/dataset.py

@@ -0,0 +1,141 @@
+from .utils import check_integrity, download_and_extract_archive
+from torch.utils.data import Dataset
+from PIL import Image
+import os
+import os.path
+import numpy as np
+import pickle
+from typing import Any, Callable, Optional, Tuple
+
+import torchvision.transforms.functional as TF
+
+
+class UpsideDownDataset(Dataset):
+    
+    """
+        Adapted from torchvision source code.
+        
+        Horizontally flips every other image and makes its label '1',
+        otherwise makes its label '0'
+    """
+    
+    base_folder = 'cifar-10-batches-py'
+    url = "https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz"
+    filename = "cifar-10-python.tar.gz"
+    tgz_md5 = 'c58f30108f718f92721af3b95e74349a'
+    train_list = [
+        ['data_batch_1', 'c99cafc152244af753f735de768cd75f'],
+        ['data_batch_2', 'd4bba439e000b95fd0a9bffe97cbabec'],
+        ['data_batch_3', '54ebc095f3ab1f0389bbae665268c751'],
+        ['data_batch_4', '634d18415352ddfa80567beed471001a'],
+        ['data_batch_5', '482c414d41f54cd18b22e5b47cb7c3cb'],
+    ]
+
+    test_list = [
+        ['test_batch', '40351d587109b95175f43aff81a1287e'],
+    ]
+    meta = {
+        'filename': 'batches.meta',
+        'key': 'label_names',
+        'md5': '5ff9c542aee3614f3951f8cda6e48888',
+    }
+
+    def __init__(
+            self,
+            root: str,
+            train: bool = True,
+            transform: Optional[Callable] = None,
+            target_transform: Optional[Callable] = None,
+            download: bool = False,
+    ) -> None:
+        
+#         super(CIFAR10, self).__init__(root, transform=transform,
+#                                       target_transform=target_transform)
+
+        self.train = train  # training set or test set
+        self.root = root
+        self.transform = transform
+        self.target_transform = target_transform
+
+        if download:
+            self.download()
+
+#         if not self._check_integrity():
+#             raise RuntimeError('Dataset not found or corrupted.' +
+#                                ' You can use download=True to download it')
+
+        if self.train:
+            downloaded_list = self.train_list
+        else:
+            downloaded_list = self.test_list
+
+        self.data: Any = []
+        self.targets = []
+
+        # now load the picked numpy arrays
+        for file_name, checksum in downloaded_list:
+            file_path = os.path.join(self.root, self.base_folder, file_name)
+            with open(file_path, 'rb') as f:
+                entry = pickle.load(f, encoding='latin1')
+                self.data.append(entry['data'])
+                if 'labels' in entry:
+                    self.targets.extend(entry['labels'])
+                else:
+                    self.targets.extend(entry['fine_labels'])
+
+        self.data = np.vstack(self.data).reshape(-1, 3, 32, 32)
+        self.data = self.data.transpose((0, 2, 3, 1))  # convert to HWC
+
+#         self._load_meta()
+        
+#     def _load_meta(self) -> None:
+#         path = os.path.join(self.root, self.base_folder, self.meta['filename'])
+#         if not check_integrity(path, self.meta['md5']):
+#             raise RuntimeError('Dataset metadata file not found or corrupted.' +
+#                                ' You can use download=True to download it')
+#         with open(path, 'rb') as infile:
+#             data = pickle.load(infile, encoding='latin1')
+#             self.classes = data[self.meta['key']]
+#         self.class_to_idx = {_class: i for i, _class in enumerate(self.classes)}
+
+    def __getitem__(self, index: int) -> Tuple[Any, Any]:
+        """
+        Args:
+            index (int): Index
+
+        Returns:
+            tuple: (image, target) where target is index of the target class.
+        """
+        img, target = self.data[index], self.targets[index]
+
+        # doing this so that it is consistent with all other datasets
+        # to return a PIL Image
+        img = Image.fromarray(img)
+        
+        if index % 2 == 0:
+            img = TF.vflip(img)
+            target = 1
+            
+        if index % 2 != 0:
+            target = 0
+            
+        if self.transform is not None:
+            img = self.transform(img)
+
+        if self.target_transform is not None:
+            target = self.target_transform(target)
+
+        return img, target
+
+
+    def __len__(self) -> int:
+        return len(self.data)
+    
+    def download(self) -> None:
+#         if self._check_integrity():
+#             print('Files already downloaded and verified')
+#             return
+        download_and_extract_archive(self.url, self.root, filename=self.filename, md5=self.tgz_md5)
+
+    def extra_repr(self) -> str:
+        return "Split: {}".format("Train" if self.train is True else "Test")

tools/utils.py

@@ -0,0 +1,470 @@
+import bz2
+import gzip
+import hashlib
+import itertools
+import lzma
+import os
+import os.path
+import pathlib
+import re
+import sys
+import tarfile
+import urllib
+import urllib.error
+import urllib.request
+import warnings
+import zipfile
+from typing import Any, Callable, List, Iterable, Optional, TypeVar, Dict, IO, Tuple, Iterator
+from urllib.parse import urlparse
+
+import requests
+import torch
+from torch.utils.model_zoo import tqdm
+
+# from .._internally_replaced_utils import (
+#     _download_file_from_remote_location,
+#     _is_remote_location_available,
+# )
+
+def _download_file_from_remote_location(fpath: str, url: str) -> None:
+    pass
+
+
+def _is_remote_location_available() -> bool:
+    return False
+
+USER_AGENT = "pytorch/vision"
+
+
+def _save_response_content(
+    content: Iterator[bytes],
+    destination: str,
+    length: Optional[int] = None,
+) -> None:
+    with open(destination, "wb") as fh, tqdm(total=length) as pbar:
+        for chunk in content:
+            # filter out keep-alive new chunks
+            if not chunk:
+                continue
+
+            fh.write(chunk)
+            pbar.update(len(chunk))
+
+
+def _urlretrieve(url: str, filename: str, chunk_size: int = 1024 * 32) -> None:
+    with urllib.request.urlopen(urllib.request.Request(url, headers={"User-Agent": USER_AGENT})) as response:
+        _save_response_content(iter(lambda: response.read(chunk_size), b""), filename, length=response.length)
+
+
+def gen_bar_updater() -> Callable[[int, int, int], None]:
+    warnings.warn("The function `gen_bar_update` is deprecated since 0.13 and will be removed in 0.15.")
+    pbar = tqdm(total=None)
+
+    def bar_update(count, block_size, total_size):
+        if pbar.total is None and total_size:
+            pbar.total = total_size
+        progress_bytes = count * block_size
+        pbar.update(progress_bytes - pbar.n)
+
+    return bar_update
+
+
+def calculate_md5(fpath: str, chunk_size: int = 1024 * 1024) -> str:
+    # Setting the `usedforsecurity` flag does not change anything about the functionality, but indicates that we are
+    # not using the MD5 checksum for cryptography. This enables its usage in restricted environments like FIPS. Without
+    # it torchvision.datasets is unusable in these environments since we perform a MD5 check everywhere.
+    md5 = hashlib.md5(**dict(usedforsecurity=False) if sys.version_info >= (3, 9) else dict())
+    with open(fpath, "rb") as f:
+        for chunk in iter(lambda: f.read(chunk_size), b""):
+            md5.update(chunk)
+    return md5.hexdigest()
+
+
+def check_md5(fpath: str, md5: str, **kwargs: Any) -> bool:
+    return md5 == calculate_md5(fpath, **kwargs)
+
+
+def check_integrity(fpath: str, md5: Optional[str] = None) -> bool:
+    if not os.path.isfile(fpath):
+        return False
+    if md5 is None:
+        return True
+    return check_md5(fpath, md5)
+
+
+def _get_redirect_url(url: str, max_hops: int = 3) -> str:
+    initial_url = url
+    headers = {"Method": "HEAD", "User-Agent": USER_AGENT}
+
+    for _ in range(max_hops + 1):
+        with urllib.request.urlopen(urllib.request.Request(url, headers=headers)) as response:
+            if response.url == url or response.url is None:
+                return url
+
+            url = response.url
+    else:
+        raise RecursionError(
+            f"Request to {initial_url} exceeded {max_hops} redirects. The last redirect points to {url}."
+        )
+
+
+def _get_google_drive_file_id(url: str) -> Optional[str]:
+    parts = urlparse(url)
+
+    if re.match(r"(drive|docs)[.]google[.]com", parts.netloc) is None:
+        return None
+
+    match = re.match(r"/file/d/(?P<id>[^/]*)", parts.path)
+    if match is None:
+        return None
+
+    return match.group("id")
+
+
+def download_url(
+    url: str, root: str, filename: Optional[str] = None, md5: Optional[str] = None, max_redirect_hops: int = 3
+) -> None:
+    """Download a file from a url and place it in root.
+
+    Args:
+        url (str): URL to download file from
+        root (str): Directory to place downloaded file in
+        filename (str, optional): Name to save the file under. If None, use the basename of the URL
+        md5 (str, optional): MD5 checksum of the download. If None, do not check
+        max_redirect_hops (int, optional): Maximum number of redirect hops allowed
+    """
+    root = os.path.expanduser(root)
+    if not filename:
+        filename = os.path.basename(url)
+    fpath = os.path.join(root, filename)
+
+    os.makedirs(root, exist_ok=True)
+
+    # check if file is already present locally
+    if check_integrity(fpath, md5):
+        print("Using downloaded and verified file: " + fpath)
+        return
+
+    if _is_remote_location_available():
+        _download_file_from_remote_location(fpath, url)
+    else:
+        # expand redirect chain if needed
+        url = _get_redirect_url(url, max_hops=max_redirect_hops)
+
+        # check if file is located on Google Drive
+        file_id = _get_google_drive_file_id(url)
+        if file_id is not None:
+            return download_file_from_google_drive(file_id, root, filename, md5)
+
+        # download the file
+        try:
+            print("Downloading " + url + " to " + fpath)
+            _urlretrieve(url, fpath)
+        except (urllib.error.URLError, OSError) as e:  # type: ignore[attr-defined]
+            if url[:5] == "https":
+                url = url.replace("https:", "http:")
+                print("Failed download. Trying https -> http instead. Downloading " + url + " to " + fpath)
+                _urlretrieve(url, fpath)
+            else:
+                raise e
+
+    # check integrity of downloaded file
+    if not check_integrity(fpath, md5):
+        raise RuntimeError("File not found or corrupted.")
+
+
+def list_dir(root: str, prefix: bool = False) -> List[str]:
+    """List all directories at a given root
+
+    Args:
+        root (str): Path to directory whose folders need to be listed
+        prefix (bool, optional): If true, prepends the path to each result, otherwise
+            only returns the name of the directories found
+    """
+    root = os.path.expanduser(root)
+    directories = [p for p in os.listdir(root) if os.path.isdir(os.path.join(root, p))]
+    if prefix is True:
+        directories = [os.path.join(root, d) for d in directories]
+    return directories
+
+
+def list_files(root: str, suffix: str, prefix: bool = False) -> List[str]:
+    """List all files ending with a suffix at a given root
+
+    Args:
+        root (str): Path to directory whose folders need to be listed
+        suffix (str or tuple): Suffix of the files to match, e.g. '.png' or ('.jpg', '.png').
+            It uses the Python "str.endswith" method and is passed directly
+        prefix (bool, optional): If true, prepends the path to each result, otherwise
+            only returns the name of the files found
+    """
+    root = os.path.expanduser(root)
+    files = [p for p in os.listdir(root) if os.path.isfile(os.path.join(root, p)) and p.endswith(suffix)]
+    if prefix is True:
+        files = [os.path.join(root, d) for d in files]
+    return files
+
+
+def _extract_gdrive_api_response(response, chunk_size: int = 32 * 1024) -> Tuple[bytes, Iterator[bytes]]:
+    content = response.iter_content(chunk_size)
+    first_chunk = None
+    # filter out keep-alive new chunks
+    while not first_chunk:
+        first_chunk = next(content)
+    content = itertools.chain([first_chunk], content)
+
+    try:
+        match = re.search("<title>Google Drive - (?P<api_response>.+?)</title>", first_chunk.decode())
+        api_response = match["api_response"] if match is not None else None
+    except UnicodeDecodeError:
+        api_response = None
+    return api_response, content
+
+
+def download_file_from_google_drive(file_id: str, root: str, filename: Optional[str] = None, md5: Optional[str] = None):
+    """Download a Google Drive file from  and place it in root.
+
+    Args:
+        file_id (str): id of file to be downloaded
+        root (str): Directory to place downloaded file in
+        filename (str, optional): Name to save the file under. If None, use the id of the file.
+        md5 (str, optional): MD5 checksum of the download. If None, do not check
+    """
+    # Based on https://stackoverflow.com/questions/38511444/python-download-files-from-google-drive-using-url
+
+    root = os.path.expanduser(root)
+    if not filename:
+        filename = file_id
+    fpath = os.path.join(root, filename)
+
+    os.makedirs(root, exist_ok=True)
+
+    if check_integrity(fpath, md5):
+        print(f"Using downloaded {'and verified ' if md5 else ''}file: {fpath}")
+
+    url = "https://drive.google.com/uc"
+    params = dict(id=file_id, export="download")
+    with requests.Session() as session:
+        response = session.get(url, params=params, stream=True)
+
+        for key, value in response.cookies.items():
+            if key.startswith("download_warning"):
+                token = value
+                break
+        else:
+            api_response, content = _extract_gdrive_api_response(response)
+            token = "t" if api_response == "Virus scan warning" else None
+
+        if token is not None:
+            response = session.get(url, params=dict(params, confirm=token), stream=True)
+            api_response, content = _extract_gdrive_api_response(response)
+
+        if api_response == "Quota exceeded":
+            raise RuntimeError(
+                f"The daily quota of the file {filename} is exceeded and it "
+                f"can't be downloaded. This is a limitation of Google Drive "
+                f"and can only be overcome by trying again later."
+            )
+
+        _save_response_content(content, fpath)
+
+
+def _extract_tar(from_path: str, to_path: str, compression: Optional[str]) -> None:
+    with tarfile.open(from_path, f"r:{compression[1:]}" if compression else "r") as tar:
+        tar.extractall(to_path)
+
+
+_ZIP_COMPRESSION_MAP: Dict[str, int] = {
+    ".bz2": zipfile.ZIP_BZIP2,
+    ".xz": zipfile.ZIP_LZMA,
+}
+
+
+def _extract_zip(from_path: str, to_path: str, compression: Optional[str]) -> None:
+    with zipfile.ZipFile(
+        from_path, "r", compression=_ZIP_COMPRESSION_MAP[compression] if compression else zipfile.ZIP_STORED
+    ) as zip:
+        zip.extractall(to_path)
+
+
+_ARCHIVE_EXTRACTORS: Dict[str, Callable[[str, str, Optional[str]], None]] = {
+    ".tar": _extract_tar,
+    ".zip": _extract_zip,
+}
+_COMPRESSED_FILE_OPENERS: Dict[str, Callable[..., IO]] = {
+    ".bz2": bz2.open,
+    ".gz": gzip.open,
+    ".xz": lzma.open,
+}
+_FILE_TYPE_ALIASES: Dict[str, Tuple[Optional[str], Optional[str]]] = {
+    ".tbz": (".tar", ".bz2"),
+    ".tbz2": (".tar", ".bz2"),
+    ".tgz": (".tar", ".gz"),
+}
+
+
+def _detect_file_type(file: str) -> Tuple[str, Optional[str], Optional[str]]:
+    """Detect the archive type and/or compression of a file.
+
+    Args:
+        file (str): the filename
+
+    Returns:
+        (tuple): tuple of suffix, archive type, and compression
+
+    Raises:
+        RuntimeError: if file has no suffix or suffix is not supported
+    """
+    suffixes = pathlib.Path(file).suffixes
+    if not suffixes:
+        raise RuntimeError(
+            f"File '{file}' has no suffixes that could be used to detect the archive type and compression."
+        )
+    suffix = suffixes[-1]
+
+    # check if the suffix is a known alias
+    if suffix in _FILE_TYPE_ALIASES:
+        return (suffix, *_FILE_TYPE_ALIASES[suffix])
+
+    # check if the suffix is an archive type
+    if suffix in _ARCHIVE_EXTRACTORS:
+        return suffix, suffix, None
+
+    # check if the suffix is a compression
+    if suffix in _COMPRESSED_FILE_OPENERS:
+        # check for suffix hierarchy
+        if len(suffixes) > 1:
+            suffix2 = suffixes[-2]
+
+            # check if the suffix2 is an archive type
+            if suffix2 in _ARCHIVE_EXTRACTORS:
+                return suffix2 + suffix, suffix2, suffix
+
+        return suffix, None, suffix
+
+    valid_suffixes = sorted(set(_FILE_TYPE_ALIASES) | set(_ARCHIVE_EXTRACTORS) | set(_COMPRESSED_FILE_OPENERS))
+    raise RuntimeError(f"Unknown compression or archive type: '{suffix}'.\nKnown suffixes are: '{valid_suffixes}'.")
+
+
+def _decompress(from_path: str, to_path: Optional[str] = None, remove_finished: bool = False) -> str:
+    r"""Decompress a file.
+
+    The compression is automatically detected from the file name.
+
+    Args:
+        from_path (str): Path to the file to be decompressed.
+        to_path (str): Path to the decompressed file. If omitted, ``from_path`` without compression extension is used.
+        remove_finished (bool): If ``True``, remove the file after the extraction.
+
+    Returns:
+        (str): Path to the decompressed file.
+    """
+    suffix, archive_type, compression = _detect_file_type(from_path)
+    if not compression:
+        raise RuntimeError(f"Couldn't detect a compression from suffix {suffix}.")
+
+    if to_path is None:
+        to_path = from_path.replace(suffix, archive_type if archive_type is not None else "")
+
+    # We don't need to check for a missing key here, since this was already done in _detect_file_type()
+    compressed_file_opener = _COMPRESSED_FILE_OPENERS[compression]
+
+    with compressed_file_opener(from_path, "rb") as rfh, open(to_path, "wb") as wfh:
+        wfh.write(rfh.read())
+
+    if remove_finished:
+        os.remove(from_path)
+
+    return to_path
+
+
+def extract_archive(from_path: str, to_path: Optional[str] = None, remove_finished: bool = False) -> str:
+    """Extract an archive.
+
+    The archive type and a possible compression is automatically detected from the file name. If the file is compressed
+    but not an archive the call is dispatched to :func:`decompress`.
+
+    Args:
+        from_path (str): Path to the file to be extracted.
+        to_path (str): Path to the directory the file will be extracted to. If omitted, the directory of the file is
+            used.
+        remove_finished (bool): If ``True``, remove the file after the extraction.
+
+    Returns:
+        (str): Path to the directory the file was extracted to.
+    """
+    if to_path is None:
+        to_path = os.path.dirname(from_path)
+
+    suffix, archive_type, compression = _detect_file_type(from_path)
+    if not archive_type:
+        return _decompress(
+            from_path,
+            os.path.join(to_path, os.path.basename(from_path).replace(suffix, "")),
+            remove_finished=remove_finished,
+        )
+
+    # We don't need to check for a missing key here, since this was already done in _detect_file_type()
+    extractor = _ARCHIVE_EXTRACTORS[archive_type]
+
+    extractor(from_path, to_path, compression)
+    if remove_finished:
+        os.remove(from_path)
+
+    return to_path
+
+
+def download_and_extract_archive(
+    url: str,
+    download_root: str,
+    extract_root: Optional[str] = None,
+    filename: Optional[str] = None,
+    md5: Optional[str] = None,
+    remove_finished: bool = False,
+) -> None:
+    download_root = os.path.expanduser(download_root)
+    if extract_root is None:
+        extract_root = download_root
+    if not filename:
+        filename = os.path.basename(url)
+
+    download_url(url, download_root, filename, md5)
+
+    archive = os.path.join(download_root, filename)
+    print(f"Extracting {archive} to {extract_root}")
+    extract_archive(archive, extract_root, remove_finished)
+
+
+def iterable_to_str(iterable: Iterable) -> str:
+    return "'" + "', '".join([str(item) for item in iterable]) + "'"
+
+
+T = TypeVar("T", str, bytes)
+
+
+def verify_str_arg(
+    value: T,
+    arg: Optional[str] = None,
+    valid_values: Iterable[T] = None,
+    custom_msg: Optional[str] = None,
+) -> T:
+    if not isinstance(value, torch._six.string_classes):
+        if arg is None:
+            msg = "Expected type str, but got type {type}."
+        else:
+            msg = "Expected type str for argument {arg}, but got type {type}."
+        msg = msg.format(type=type(value), arg=arg)
+        raise ValueError(msg)
+
+    if valid_values is None:
+        return value
+
+    if value not in valid_values:
+        if custom_msg is not None:
+            msg = custom_msg
+        else:
+            msg = "Unknown value '{value}' for argument {arg}. Valid values are {{{valid_values}}}."
+            msg = msg.format(value=value, arg=arg, valid_values=iterable_to_str(valid_values))
+        raise ValueError(msg)
+
+    return value

train.py

@@ -0,0 +1,26 @@
+import torch
+
+def train(model, trainloader, optimizer, criterion, DEVICE):
+    model.train()
+    
+    running_loss = 0
+    for itr, data in enumerate(trainloader):
+#         print(itr)
+#         print(data[0].shape, data[1].shape)
+#         print(len(trainloader))
+#         if itr % 100 == 0:
+#             print("itr: {}".format(itr))
+        optimizer.zero_grad()
+        
+        imgs, target = data[0].to(DEVICE), data[1].to(DEVICE)
+        output_logits = model(imgs)
+        loss = criterion( output_logits, target)
+        
+        running_loss = loss.item()
+        loss.backward()
+        optimizer.step()
+        
+    epoch_loss = running_loss/len(trainloader)
+    print("epoch loss = {}".format(epoch_loss))
+
+    return epoch_loss

viz/save_plot.py

@@ -0,0 +1,6 @@
+import numpy as np
+import matplotlib.pyplot as plt
+
+def save_plot(*arrays):
+    for array in args:
+