Skip to content

pytorch: Open-source machine learning library

Overview of package

Package: pytorch
Description: Open-source machine learning library
For more information: https://pytorch.org/
Categories: ComputerVision
MachineLearning
NumericalMethods
PythonModule
License: OpenSource (BSD)

General Usage

PyTorch is an open-source machine learning library that evolved from the (no longer supported) Lua based Torch library. It is commonly used for computer vision and natural language processing tasks. Among its high level features are:

  • Tensor computing (like NumPy) with strong GPU acceleration
  • Deep neural networks built on a tape-based automatic differentiation system

This module will add the torch, torchvision, and torchsummary modules to your PYTHONPATH Please note that despite the module being named "pytorch", you should use "import torch" or similar in your Python code

Available versions by cluster

This section lists the available versions of the package pytorchon the UMD HPC clusters.

Versions available on the Zaratan cluster

Version Module tags CPU(s) optimized for GPU ready?
2.0.1 2.0.1 icelake, zen2 Yes
1.11.0 1.11.0 zen2 Yes

GPU suppport

The PyTorch package can make use of GPUs on nodes with GPUs. There is nothing special that needs to be done in the module load or the various pytorch* commands, but you will need to instruct the package to use the GPUs within your python code. This is typically done by replacing a line like

device = torch.device("cpu")
with something like
device = torch.device("cuda:0")

Distributed pytorch

Although the Singularity containers with pytorch do not have MPI support, pytorch has its own distributed package (torch.distributed) which can handle parallelizing your computations across multiple nodes. More information on using torch.distributed in your Python codes can be found at the PyTorch Distributed Tutorial and the Distributed Communication Documentation.

We recommend using the TCP based initialization, using something like the example script below:

#!/bin/bash
#SBATCH -ntasks=40
#SBATCH -t 00:01:00
#SBATCH --mem-per-cpu=2048
#SBATCH --exclusive

#Define module command, etc
. ~/.profile
#Load the pytorch module
module load pytorch/0.4.1

#Number of processes per node to launch (20 for CPU, 2 for GPU)
NPROC_PER_NODE=20

#This command to run your pytorch script
#You will want to replace this
COMMAND="YOUR_TRAINING_SCRIPT.py --arg1 --arg2 ..."

#We want names of master and slave nodes
MASTER=`/bin/hostname -s`
SLAVES=`scontrol show hostnames $SLURM_JOB_NODELIST | grep -v $MASTER`
#Make sure this node (MASTER) comes first
HOSTLIST="$MASTER $SLAVES"

#Get a random unused port on this host(MASTER) between 2000 and 9999
#First line gets list of unused ports
#2nd line restricts between 2000 and 9999
#3rd line gets single random port from the list
MPORT=`ss -tan | awk '{print $4}' | cut -d':' -f2 | \
        grep "[2-9][0-9]\{3,3\}" | grep -v "[0-9]\{5,5\}" | \
        sort | uniq | shuf`

#Launch the pytorch processes, first on master (first in $HOSTLIST) then
#on the slaves
RANK=0
for node in $HOSTLIST; do
        ssh -q $node \
                pytorch -m torch.distributed.launch \
                --nproces_per_node=$NPROCS_PER_NODE \
                --nnodes=$SLURM_JOB_NUM_NODES \
                --node_rank=$RANK \
                --master_addr="$MASTER" --master_port="$MPORT" \
                $COMMAND &
        RANK=$((RANK+1))
done
wait

import argparse
import torch.distributed as dist

parser.add_argument('--distributed', action='store_true', help='enables distributed processes')
parser.add_argument('--local_rank', default=0, type=int, help='number of distributed processes')
parser.add_argument('--dist_backend', default='gloo', type=str, help='distributed backend')

def main():
opt = parser.parse_args()
if opt.distributed:
    dist.init_process_group(backend=opt.dist_backend, init_method='env://')

print("Initialized Rank:", dist.get_rank())

if __name__ == '__main__':
    main()