PyTorch is a package intended for deep learning in Python and C++. It’s especially popular for computer vision and natural language processing applications.

Quick links:

Standard import statement:

import torch  
import torch.nn as nn  
import torch.nn.functional as F  
import torch.optim as optim

Classes

PyTorch has a few important sub-classes:

  • nn contains neural network-related building blocks, including activation functions, loss functions, and linear transformations.
    • nn.functional contains activation functions.
  • optim contains optimisation algorithms.

The core data structure of PyTorch is the tensor (torch.Tensor), which functions similarly to a NumPy array.

  • To convert from a NumPy array, use torch.from_numpy(arr).

CUDA

PyTorch also has broad support for parallelisation via CUDA. By default it uses the CPU to run computations, so some quick initialisation steps are needed:

  • torch.cuda.is_available() returns True if there’s a CUDA-compatible device.
  • torch.cuda.device_count() tells us how many possible devices we can run on.
  • torch.cuda.current_device() gives us which device we’re using right now.
  • torch.cuda.get_device_name(n: int) outputs the name of the device. Use with the current_device() method.
  • torch.cuda.set_default_device('cuda' if torch.cuda.is_available() else 'cpu') to switch to CUDA.

Standard CUDA switches:

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  
print(device)  
torch.set_default_device(device)  
torch.set_default_dtype(torch.float32)

The .to(device: str) method transfers an object to CUDA memory. To make sure we’re running on CUDA, we should do the following. The best place to do this seems to be in a training function.

  • We need to transfer the model itself.
  • Each data point, when we’re loading from a data loader class while training.
  • In the data loader, the random number generator should also be on the GPU, with generator=torch.Generator(device='cuda') as an extra parameter in the torch.utils.data.DataLoader object.

Installation

Problem: CUDA-specific installs of PyTorch are a little difficult to get right. This detects and wraps the right pip commands.

pip install light-the-torch
ltt install torch

See also