Custom transforms in PyTorch

Introduction

Custom transforms let you change your data in your own way before training. This helps your model learn better by seeing different versions of the data.

You want to flip or rotate images in a special way not covered by built-in tools.

You need to add noise or change colors in a unique style to make your model stronger.

Your data is not images but text or numbers, and you want to prepare it before training.

You want to combine several small changes into one easy step.

You want to fix or clean your data in a way no standard tool does.

Syntax

PyTorch

class MyTransform:
    def __init__(self, param):
        self.param = param

    def __call__(self, sample):
        # change sample using self.param
        transformed_sample = sample  # example placeholder
        return transformed_sample

The class must have a __call__ method to work like a function.

You can add parameters in __init__ to control how the transform works.

Examples

This transform adds random noise to a tensor sample.

PyTorch

import torch

class AddNoise:
    def __init__(self, noise_level):
        self.noise_level = noise_level

    def __call__(self, sample):
        return sample + self.noise_level * torch.randn_like(sample)

This transform converts a color image tensor to grayscale by averaging color channels.

PyTorch

class ToGray:
    def __call__(self, image):
        return image.mean(dim=0, keepdim=True)

Sample Model

This example shows a simple custom transform that doubles the input tensor values.

PyTorch

import torch
from torchvision import transforms

class MultiplyByTwo:
    def __call__(self, x):
        return x * 2

# Create a tensor sample
sample = torch.tensor([1.0, 2.0, 3.0])

# Create transform instance
transform = MultiplyByTwo()

# Apply transform
transformed_sample = transform(sample)

print('Original:', sample)
print('Transformed:', transformed_sample)

Important Notes

Custom transforms can be combined with built-in transforms using transforms.Compose.

Make sure your transform works with the data type you use (e.g., PIL images or tensors).

Test your transform separately to avoid bugs in your training pipeline.

Summary

Custom transforms let you change data your own way before training.

They are classes with a __call__ method that changes the input.

Use them to add new data changes not in standard libraries.