#!/bin/bash

# assumes `cuda-repo-ubuntu1604_8.0.61-1_amd64.deb` and `cudnn-8.0-linux-x64-v6.0.tgz` are located in working directory

# download `cuda-repo-ubuntu1604_8.0.61-1_amd64.deb` here: https://developer.nvidia.com/cuda-downloads (linux -> x86_64 -> Ubuntu -> 16.04 -> deb (network))
# download `cudnn-8.0-linux-x64-v6.0.tgz` here: https://developer.nvidia.com/cudnn

# install CUDA Toolkit v8.0
# instructions from https://developer.nvidia.com/cuda-downloads (linux -> x86_64 -> Ubuntu -> 16.04 -> deb (network))
sudo dpkg -i cuda-repo-ubuntu1604_8.0.61-1_amd64.deb
sudo apt-get update
sudo apt-get install cuda

# install cuDNN v6.0
tar -xzvf cudnn-8.0-linux-x64-v6.0.tgz
sudo cp -P cuda/include/cudnn.h /usr/local/cuda-8.0/include
sudo cp -P cuda/lib64/libcudnn* /usr/local/cuda-8.0/lib64/
sudo chmod a+r /usr/local/cuda-8.0/lib64/libcudnn*

# set environment variables
export PATH=/usr/local/cuda-8.0/bin${PATH:+:${PATH}}
export LD_LIBRARY_PATH=/usr/local/cuda-8.0/lib64\${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}

# [optional] install Python 3.6
sudo add-apt-repository ppa:jonathonf/python-3.6
sudo apt-get update
sudo apt-get install -y python3.6 python3-pip

# install pytorch and it's python package dependencies
sudo pip3 install -U numpy
sudo pip3 install -U pyyaml
# Linux, Python 3.6, CUDA Toolkit 8.0
sudo pip3 install -U http://download.pytorch.org/whl/cu80/torch-0.1.11.post5-cp36-cp36m-linux_x86_64.whl
sudo pip3 install -U torchvision