Skip to content

Decorators: Framework Magic

Overview

Decorators are core to how PyTorch works: - @torch.jit.script = Compile function to TorchScript - @property = Lazy evaluation of model attributes - @abstractmethod = Define interfaces - Custom decorators = Add functionality to layers, profiling, debugging

Decorators are "functions that take functions and return functions."

Decorator Basics

Simple Decorator

def my_decorator(func):
    """Wraps a function."""
    def wrapper(*args, **kwargs):
        print(f"Before calling {func.__name__}")
        result = func(*args, **kwargs)
        print(f"After calling {func.__name__}")
        return result
    return wrapper

@my_decorator  # Equivalent to: add = my_decorator(add)
def add(a, b):
    return a + b

add(1, 2)
# Output:
# Before calling add
# After calling add

How Decorators Work

# These are equivalent:

# Using @ syntax
@my_decorator
def add(a, b):
    return a + b

# Explicit call
def add(a, b):
    return a + b
add = my_decorator(add)

PyTorch Decorators

@property: Lazy Evaluation

class Tensor:
    def __init__(self, data):
        self._data = data
        self._shape = None  # Computed lazily

    @property
    def shape(self):
        """Compute shape on demand."""
        if self._shape is None:
            self._shape = tuple(len(d) for d in self._data)
        return self._shape

t = Tensor([[1, 2], [3, 4]])
print(t.shape)  # (2, 2) - computed when accessed

@staticmethod and @classmethod

class Tensor:
    count = 0  # Class variable

    def __init__(self, data):
        self.data = data
        Tensor.count += 1

    @staticmethod
    def create_zeros(shape):
        """Static method (no self)."""
        return Tensor([0] * shape[0])

    @classmethod
    def from_numpy(cls, array):
        """Class method (has cls)."""
        return cls(array.tolist())

# Usage
t1 = Tensor.create_zeros((3,))
t2 = Tensor.from_numpy(np.array([1, 2, 3]))
print(Tensor.count)  # 2

@abstractmethod

from abc import ABC, abstractmethod

class Layer(ABC):
    @abstractmethod
    def forward(self, x):
        """Must be implemented by subclasses."""
        pass

class Linear(Layer):
    def forward(self, x):
        return x @ self.weight

# Linear()(x) works - forward implemented
# Layer()(x) fails - abstract method not implemented

Custom Decorators for ML

Timing Decorator

import functools
import time

def timing_decorator(func):
    """Measure function execution time."""
    @functools.wraps(func)  # Preserve function metadata
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        elapsed = time.time() - start
        print(f"{func.__name__} took {elapsed:.3f}s")
        return result
    return wrapper

@timing_decorator
def train_epoch(model, data):
    # Training code...
    return loss

train_epoch(model, data)
# Output: train_epoch took 12.345s

Validation Decorator

def validate_inputs(func):
    """Validate tensor inputs."""
    @functools.wraps(func)
    def wrapper(self, x):
        if not isinstance(x, torch.Tensor):
            raise TypeError(f"Expected Tensor, got {type(x)}")
        if x.dim() == 0:
            raise ValueError("Tensor cannot be 0-dimensional")
        return func(self, x)
    return wrapper

class Model(torch.nn.Module):
    @validate_inputs
    def forward(self, x):
        return self.fc(x)

model = Model(torch.nn.Linear(10, 5))
model(torch.randn(3, 10))  # OK
model(42)  # TypeError: Expected Tensor
model(torch.tensor(1))  # ValueError: Tensor cannot be 0-dimensional

Gradient Tracking Decorator

def with_grad_tracking(func):
    """Track gradients for debugging."""
    @functools.wraps(func)
    def wrapper(self, *args, **kwargs):
        results = []
        for i, arg in enumerate(args):
            if isinstance(arg, torch.Tensor) and arg.requires_grad:
                print(f"Input {i} requires grad")

        output = func(self, *args, **kwargs)

        if isinstance(output, torch.Tensor) and output.requires_grad:
            print(f"Output requires grad")

        return output
    return wrapper

class Model(torch.nn.Module):
    @with_grad_tracking
    def forward(self, x):
        return self.fc(x)

Decorators with Arguments

Fixed Arguments

def repeat_n_times(n):
    """Repeat function call n times."""
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            results = []
            for _ in range(n):
                result = func(*args, **kwargs)
                results.append(result)
            return results
        return wrapper
    return decorator

@repeat_n_times(3)
def create_tensor():
    return torch.randn(3)

tensors = create_tensor()  # Creates 3 tensors
print(len(tensors))  # 3

Variable Arguments

def retry(max_attempts=3, delay=1):
    """Retry function on failure."""
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            import time
            for attempt in range(max_attempts):
                try:
                    return func(*args, **kwargs)
                except Exception as e:
                    if attempt == max_attempts - 1:
                        raise
                    print(f"Attempt {attempt+1} failed, retrying in {delay}s...")
                    time.sleep(delay)
        return wrapper
    return decorator

@retry(max_attempts=3, delay=2)
def load_model(path):
    return torch.load(path)

model = load_model("model.pt")

Chaining Decorators

Multiple decorators applied bottom-to-top:

@timing_decorator
@validate_inputs
def forward(self, x):
    return self.fc(x)

# Equivalent to:
def forward(self, x):
    return self.fc(x)
forward = validate_inputs(forward)
forward = timing_decorator(forward)

# When called:
# 1. timing_decorator wrapper called
# 2. Inside, calls validate_inputs wrapper
# 3. Inside, calls original forward
# 4. Returns to validate_inputs wrapper, then timing_decorator wrapper

Practical: Combined Decorators

def debug(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        print(f"Calling {func.__name__}")
        return func(*args, **kwargs)
    return wrapper

def validate(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        if len(args) == 0:
            raise ValueError("Need arguments")
        return func(*args, **kwargs)
    return wrapper

@debug
@validate
def forward(x):
    return x * 2

forward(5)
# Output:
# Calling forward
# 10

forward()  # Raises ValueError before debug prints

Class Decorators

Decorator on Class

def add_repr(cls):
    """Add repr to any class."""
    def __repr__(self):
        attrs = ", ".join(f"{k}={v}" for k, v in self.__dict__.items())
        return f"{cls.__name__}({attrs})"

    cls.__repr__ = __repr__
    return cls

@add_repr
class Point:
    def __init__(self, x, y):
        self.x = x
        self.y = y

p = Point(3, 4)
print(p)  # Point(x=3, y=4)

Registering Classes

# Registry pattern (used in PyTorch)
MODEL_REGISTRY = {}

def register_model(name):
    def decorator(cls):
        MODEL_REGISTRY[name] = cls
        return cls
    return decorator

@register_model("resnet50")
class ResNet50(torch.nn.Module):
    pass

@register_model("vgg16")
class VGG16(torch.nn.Module):
    pass

# Later: Load models dynamically
model_class = MODEL_REGISTRY["resnet50"]
model = model_class()

Advanced: Functools.wraps

Preserves function metadata:

import functools

def bad_decorator(func):
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs)
    return wrapper

def good_decorator(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs)
    return wrapper

def add(x, y):
    """Add two numbers."""
    return x + y

bad = bad_decorator(add)
good = good_decorator(add)

print(bad.__name__)   # 'wrapper' (lost original)
print(good.__name__)  # 'add' (preserved)

print(bad.__doc__)   # None
print(good.__doc__)  # 'Add two numbers.'

Decorator Stacking Performance

import timeit

def timer_decorator(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs)
    return wrapper

# Single decorator
@timer_decorator
def func1(x):
    return x * 2

# Multiple decorators
@timer_decorator
@timer_decorator
@timer_decorator
def func3(x):
    return x * 2

time1 = timeit.timeit(lambda: func1(42), number=1000000)
time3 = timeit.timeit(lambda: func3(42), number=1000000)

print(f"Single: {time1:.3f}s")
print(f"Triple: {time3:.3f}s")
# Overhead is minimal (<5%)

PyTorch-Specific Decorators

TorchScript Compilation

@torch.jit.script
def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    return x + y

# Compiled to TorchScript (faster, deployable)
x = torch.tensor([1.0, 2.0])
y = torch.tensor([3.0, 4.0])
print(add(x, y))

Custom Gradient

class CustomFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        return x ** 2

    @staticmethod
    def backward(ctx, grad_output):
        x, = ctx.saved_tensors
        return 2 * x * grad_output

# Use in model
x = torch.tensor([2.0], requires_grad=True)
y = CustomFunction.apply(x)
y.backward()
print(x.grad)  # tensor([4.])

Common Pitfalls

Pitfall 1: Forgetting functools.wraps

# BAD: Loses function metadata
def bad_decorator(func):
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs)
    return wrapper

# GOOD: Preserves metadata
def good_decorator(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs)
    return wrapper

Pitfall 2: Decorators Changing Function Signature

# This hides the actual signature
def decorator(func):
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs)
    return wrapper

@decorator
def add(a, b):
    return a + b

# IDE can't show add(a, b) signature
# Instead shows wrapper(*args, **kwargs)

Summary

  • Decorators = Functions that modify functions
  • @property = Make methods look like attributes
  • @staticmethod = Function bound to class, not instance
  • @classmethod = Method receives class as first argument
  • Custom decorators = Add functionality (timing, validation, caching)
  • Chaining = Multiple decorators compose
  • functools.wraps = Always use to preserve metadata