10 Essential PyTorch Transformations for Computer Vision

PyTorch is one of the most popular deep learning frameworks used today, especially for computer vision tasks. A big reason for its widespread adoption is the extensive torchvision package which provides access to popular datasets, model architectures, and image transformations.

In computer vision, we often need to preprocess and manipulate image data before feeding it into a deep learning model. This is where PyTorch‘s transformations come in handy. Transformations allow us to perform operations like resizing, cropping, normalizing, and augmenting images with ease.

In this post, we‘ll go through 10 essential PyTorch transformations you should know for computer vision tasks. We‘ll cover basic transformations, dive into the powerful transforms.Compose() functionality, and also touch on some more advanced techniques. Let‘s get started!

1. Resizing and Cropping

One of the most fundamental transformations is resizing an image. Deep learning models typically require input images to have uniform dimensions, so resizing is often a necessary first step. You can resize an image using transforms.Resize():

from torchvision import transforms

transform = transforms.Resize(256)
resized_img = transform(img)

This will resize the shortest side of the image to 256 pixels while preserving the aspect ratio. You can also specify an exact size using a tuple:

transform = transforms.Resize((128, 256))
resized_img = transform(img)  # resizes to 128x256

In addition to resizing, cropping is frequently used to extract regions of interest from an image. transforms.CenterCrop() crops the center portion of an image:

transform = transforms.CenterCrop(128)
cropped_img = transform(img)  # takes 128x128 center crop

There are also other crop variations like transforms.RandomCrop() which takes a random crop, and the more general transforms.FiveCrop() and transforms.TenCrop() for taking multiple crops of an image.

2. Flipping and Rotating

Another set of basic transformations are flips and rotations. These can be used for data augmentation to improve model robustness. transforms.RandomHorizontalFlip() will randomly flip an image horizontally with a given probability (default 0.5):

transform = transforms.RandomHorizontalFlip(p=0.7)
flipped_img = transform(img) 

Similarly, transforms.RandomVerticalFlip() performs vertical flips. For rotations, there‘s transforms.RandomRotation() which rotates an image by a random angle within a specified range:

transform = transforms.RandomRotation(degrees=30)
rotated_img = transform(img)  # randomly rotates up to 30 degrees

3. Converting to Tensor and Normalizing

To fed an image into a PyTorch model, we first need to convert it to a tensor. This is done using transforms.ToTensor():

transform = transforms.ToTensor()
tensor_img = transform(img)

This converts a PIL image or NumPy array to a PyTorch tensor with values scaled to [0, 1]. It also switches the ordering of dimensions from HWC to CHW format which PyTorch uses.

Another important preprocessing step is normalization. This adjusts the range of pixel intensities to be centered around 0 with a standard deviation of 1, which helps with model convergence. We can normalize using transforms.Normalize():

transform = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
normalized_img = transform(tensor_img)

The mean and std values here are standard for models pretrained on ImageNet. You‘ll want to adjust them if using a different dataset.

4. Composing Transformations

While applying transformations individually works, it‘s usually cleaner and more convenient to compose them together into a single callable object. That‘s where transforms.Compose() comes in:

transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

tensor_img = transform(img)

Here we resize an image, take a center crop, convert to a tensor, and normalize the pixel values all in one go. The output tensor_img can be directly passed to a model. Being able to easily compose transformations like this is a big convenience of the PyTorch API.

5. Data Augmentation

Beyond the basic transformations, PyTorch also provides many operations geared towards data augmentation. The goal of augmentation is to synthetically expand the training set by applying random (but realistic) modifications to images. This helps combat overfitting and improves generalization. Some built-in augmentations include:

• transforms.ColorJitter(): Randomly changes brightness, contrast, saturation and hue of an image
• transforms.GaussianBlur(): Applies Gaussian blur to an image
• transforms.RandomAffine(): Performs a random affine transformation (translation, rotation, shear, scaling)
• transforms.RandomPerspective(): Applies a random perspective transformation
• transforms.RandomErasing(): Randomly erases a rectangular region of an image

These augmentations can produce an effectively infinite supply of unique training examples. As an example, here‘s how we could use ColorJitter and GaussianBlur:

transform = transforms.Compose([
    transforms.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.1),
    transforms.GaussianBlur(kernel_size=7),
    transforms.ToTensor()
])

augmented_img = transform(img)

6. Lambda Transformations

For custom transformations not covered by the built-in options, we can use transforms.Lambda to wrap an arbitrary function:

transform = transforms.Lambda(lambda x: x.add(10))

This will add 10 to each pixel value. Lambda transformations provide flexibility to perform any operation you need.

7. Conversion Transformations

There are also transformations for converting between different image formats and data types:

• transforms.ToPILImage(): Converts a tensor or NumPy array to a PIL image
• transforms.ToTensor(): Converts a PIL image or NumPy array to a tensor
• transforms.PILToTensor(): More efficient version of ToTensor() for PIL images

8. Functional Transforms

For more fine-grained control, the torchvision.transforms.functional module provides functional versions of all the transformations. Instead of being callable classes, these are just functions that take an image and any necessary parameters.

For example, to resize an image using the functional form:

from torchvision.transforms import functional as TF
resized_img = TF.resize(img, 256) 

The functional transforms can be useful if you need to apply a transformation conditionally or don‘t want to construct a callable transform object.

9. Scriptable Transformations

Starting from PyTorch 1.7, there‘s a new transforms.Compose that‘s compatible with TorchScript compilation. This allows for transformations to be serialized as part of a model and later run in a non-Python environment like C++.

Using scriptable transforms is as easy as replacing transforms.Compose with torch.nn.Sequential:

transform = torch.nn.Sequential(
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ConvertImageDtype(torch.float),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
)

The scriptable transforms have a few constraints (e.g. no Lambdas) but cover most common use cases.

10. Optimizing Transformations

A final tip is to be thoughtful about where transformations are applied. The earlier you can perform transformations in your data pipeline, the better. Transformations done in a DataLoader‘s worker processes can needlessly bottleneck training.

When possible, try to perform transformations ahead of time and serialize the results to disk. For example, when working with a custom dataset, you could apply transforms.Resize(), transforms.CenterCrop() and transforms.ToTensor() to the source images and save the resulting tensors. Then at training time, you only need to load the preprocessed tensors and apply normalization.

Avoiding redundant transformations can significantly speed up data loading and training. See PyTorch‘s performance tutorial for more tips.

Conclusion

To recap, we covered 10 essential PyTorch transformations for computer vision:

1. Resizing and cropping
2. Flipping and rotating
3. Converting to tensor and normalizing
4. Composing transformations
5. Data augmentation
6. Lambda transformations
7. Conversion transformations
8. Functional transforms
9. Scriptable transforms
10. Optimizing transformations

Transformations are a core part of working with image data in PyTorch. Knowing how and when to use them will make your computer vision workflows much more effective. Hopefully this post has given you a solid starting point. For more details, be sure to check out the official torchvision documentation.

Thanks for reading! Let me know if you have any other handy PyTorch transformations that I missed.