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:
- Documentation (stable version)
- GitHub repository
- Related packages
Standard import statement:
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()
returnsTrue
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 thecurrent_device()
method.torch.cuda.set_default_device('cuda' if torch.cuda.is_available() else 'cpu')
to switch to CUDA.
Standard CUDA switches:
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 thetorch.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.