PYTORCH CHEAT SHEET

Imports
General

import torch # root package

from torch.utils.data import Dataset, DataLoader # dataset representation and loading

Neural Network API

import torch.autograd as autograd # computation graph

from torch import Tensor # tensor node in the computation graph

import torch.nn as nn # neural networks

import torch.optim as optim # optimizers e.g. gradient descent, ADAM, etc.

from torch.jit import script, trace # hybrid frontend decorator and tracing jit

See autograd, nn, functional and optim

Torchscript and JIT

torch.jit.trace() # takes your module or function and an example

# data input, and traces the computational steps

# that the data encounters as it progresses through the model

@script # decorator used to indicate data-dependent

# control flow within the code being traced

See Torchscript

ONNX

torch.onnx.export(model, dummy data, xxxx.proto) # exports an ONNX formatted

# model using a trained model, dummy

# data and the desired file name

model = onnx.load("alexnet.proto") # load an ONNX model

onnx.checker.check_model(model) # check that the model

# IR is well formed

onnx.helper.printable_graph(model.graph) # print a human readable

# representation of the graph

See onnx

Vision

from torchvision import datasets, models, transforms # vision datasets,

# architectures &

# transforms

import torchvision.transforms as transforms # composable transforms

See torchvision

Distributed Training

import torch.distributed as dist # distributed communication

from torch.multiprocessing import Process # memory sharing processes

See distributed and multiprocessing

Tensors
https://pytorch.org/tutorials/beginner/ptcheat.html 1/3
Creation
 torch.cuda.is_available # check for cuda

x = x.cuda() # move x's data from

# CPU to GPU and return new object

x = x.cpu() # move x's data from GPU to CPU

# and return new object

if not args.disable_cuda and torch.cuda.is_available(): # device agnostic code

args.device = torch.device('cuda') # and modularity

else: #

args.device = torch.device('cpu') #

net.to(device) # recursively convert their

# parameters and buffers to

# device specific tensors

x = x.to(device) # copy your tensors to a device

# (gpu, cpu)

See cuda

Deep Learning

nn.Linear(m,n) # fully connected layer from

# m to n units

nn.ConvXd(m,n,s) # X dimensional conv layer from

# m to n channels where X⍷{1,2,3}

# and the kernel size is s

nn.MaxPoolXd(s) # X dimension pooling layer

# (notation as above)

nn.BatchNormXd # batch norm layer

nn.RNN/LSTM/GRU # recurrent layers

nn.Dropout(p=0.5, inplace=False) # dropout layer for any dimensional input

nn.Dropout2d(p=0.5, inplace=False) # 2-dimensional channel-wise dropout

nn.Embedding(num_embeddings, embedding_dim) # (tensor-wise) mapping from

# indices to embedding vectors

See nn

Loss Functions

nn.X # where X is L1Loss, MSELoss, CrossEntropyLoss

# CTCLoss, NLLLoss, PoissonNLLLoss,

# KLDivLoss, BCELoss, BCEWithLogitsLoss,

# MarginRankingLoss, HingeEmbeddingLoss,

# MultiLabelMarginLoss, SmoothL1Loss,

# SoftMarginLoss, MultiLabelSoftMarginLoss,

# CosineEmbeddingLoss, MultiMarginLoss,
# or TripletMarginLoss

See loss functions

Activation Functions

nn.X # where X is ReLU, ReLU6, ELU, SELU, PReLU, LeakyReLU,

# RReLu, CELU, GELU, Threshold, Hardshrink, HardTanh,

# Sigmoid, LogSigmoid, Softplus, SoftShrink,

# Softsign, Tanh, TanhShrink, Softmin, Softmax,

# Softmax2d, LogSoftmax or AdaptiveSoftmaxWithLoss

See activation functions

Optimizers

opt = optim.x(model.parameters(), ...) # create optimizer

opt.step() # update weights

optim.X # where X is SGD, Adadelta, Adagrad, Adam,

# AdamW, SparseAdam, Adamax, ASGD,

# LBFGS, RMSprop or Rprop

https://pytorch.org/tutorials/beginner/ptcheat.html 2/3
See optimizers