Running PyTorch

Running PyTorch on CPU

Install pre-built PyTorch from Anaconda.

$ module load miniconda3/4.5
$ conda config --add channels https://mirrors.sjtug.sjtu.edu.cn/anaconda/pkgs/free/
$ conda config --add channels https://mirrors.sjtug.sjtu.edu.cn/anaconda/pkgs/main/
$ conda config --set show_channel_urls yes
$ conda create --name torch-py3-cpu numpy ipython jupyter 
$ source activate torch-py3-cpu
$ conda install pytorch-cpu torchvision -c pytorch

Verify the installation by loading PyTorch module.DO NOT run TenorFlow jobs on login nodes.

$ python -c 'import torch; print(torch.__version__)'
0.4.0

Log out and log in again, then request an interactive SLURM job to run PyTorch.

$ module purge
$ srun -p k80 -N 1 --exclusive --pty /bin/bash
hostname
module load miniconda3/4.5
source activate torch-py3-cpu
python -c 'import torch; print(torch.__version__)'

Prepare a job script torchcpu.slurm with contents as follows.

#!/bin/bash

#SBATCH -J torch-cpu
#SBATCH -p k80
#SBATCH --mail-type=end
#SBATCH --mail-user=YOU@EMAIL.COM
#SBATCH -o %j.out
#SBATCH -e %j.err
#SBATCH -n 1
#SBATCH --exclusive

source /usr/share/Modules/init/bash
module purge
module load miniconda3/4.5

source activate torch-py3-cpu
python -c 'import torch; print(torch.__version__)'

Then submit this job to SLURM. Please refer to https://pi.sjtu.edu.cn/doc/slurm for SLURM usage.

$ sbatch torchcpu.slurm

Running PyTorch on GPU

Install pre-built GPU PyTorch from Anaconda.

$ module load miniconda3/4.5
$ conda config --add channels https://mirrors.sjtug.sjtu.edu.cn/anaconda/pkgs/free/
$ conda config --add channels https://mirrors.sjtug.sjtu.edu.cn/anaconda/pkgs/main/
$ conda config --set show_channel_urls yes
$ conda create --name torch-py3-gpu numpy ipython jupyter
$ source activate torch-py3-gpu
$ conda install pytorch torchvision cuda91 -c pytorch

Verify the installation by loading PyTorch module on mu07(202.120.58.231).DO NOT run TenorFlow jobs on login nodes.

$ module purge; module load miniconda3/4.5 cuda/9.0 cudnn/7.0
$ source activate torch-py3-gpu
$ python -c 'import torch; print(torch.__version__)'
0.4.0

Log out and log again, then request an interactive SLURM job to run PyTorch.

$ module purge
$ srun -p k80 -N 1 --exclusive --gres=gpu:2 --pty /bin/bash
hostname
module load miniconda3/4.5 cuda/9.0 cudnn/7.0
source activate torch-py3-gpu
python -c 'import torch; print(torch.__version__)'
0.4.0

Prepare a job script torchgpu.slurm with contents as follows.

#!/bin/bash

#SBATCH -J torch-gpu
#SBATCH -p k80
#SBATCH --mail-type=end
#SBATCH --mail-user=YOU@EMAIL.COM
#SBATCH -o %j.out
#SBATCH -e %j.err
#SBATCH -N 1
#SBATCH --exclusive

source /usr/share/Modules/init/bash
module purge
module load miniconda3/4.5 cuda/9.0 cudnn/7.0

source activate torch-py3-gpu
python -c 'import torch; print(torch.__version__)'

Then submit this job to SLURM. Please refer to https://pi.sjtu.edu.cn/doc/slurm for SLURM usage.

$ sbatch torchgpu.slurm

Reference