Speed Up Your PyTorch Development Using Types
Tags: machinelearning, python
Published: 2024-01-20 20:28:00

Types in Python are optional, but for any PyTorch project exceeding 1,000 lines of code, implementing them can be highly beneficial.

You've probably encountered a scenario where, after coding a model with several networks and initiating training, you realize there's a mismatch in data dimensions or type. Types help clarify the information flow. Moreover, they enable code completion tools to suggest more appropriate options.

While types won't eliminate debugging, they do make the process more focused.

## Python Types

First, check if your editor or IDE supports static type checking. If not, a quick online search will reveal popular options like `Mypy` or `Pyright`.

Python types are akin to a simplified version of the typing systems in Java or C++.

Here's how the syntax looks:
```python
variable: int = 10
```
For functions, specify the types of parameters and the return type:
```python
def function(x: float, y: float) -> float:
    return x ** y
```

Below is a list of the most fundamental Python types you'll encounter:

```python
int    # Whole number integers
float  # Floating point numbers
bool   # Boolean values (True or False)
str    # Strings
list   # Ordered mutable lists
tuple  # Ordered immutable lists
dict   # Dictionaries of key-value pairs
```

For the latter three, use the Python typing package for more precise type specifications:
```python
from typing import Dict, List, Tuple
```
Example usages:
```python
a_dict: Dict[str, float] = {'a': 1.0, ...}
a_list: List[int] = [1, 2, 3]
a_tuple: Tuple[int, bool, float] = (1, True, 3.141)

def forward(input: List[float]) -> Dict[str, float]:
    ...
```

For dealing with uncertain types, `Any` is a versatile type that accepts any value. Use it sparingly to maintain the effectiveness of your type checker.

The `Self` type represents the current class:
```python
import numpy as np
from typing import Self, List, Tuple

@dataclass
class PointCloud:
    points: List[Tuple[float, float, float]]

    def from_numpy(arr: np.ndarray) -> Self:
        ...
```

Frequently recurring patterns in your codebase may warrant the creation of custom types from basic types:
```python
from typing import Dict, List

SomeType = Dict[str, List[int]]
var: SomeType = ...
```
However, using expressive classes as types is often more sensible:
```python
pcd: PointCloud = PointCloud([(1.0, 1.0, 1.0), ...])  # example from above
```
Be aware that your static type checker can often automatically identify the correct class upon assignment. However, use explicit types in scenarios where automatic deduction is not possible.

## Types for PyTorch

A clean architecture combined with proper types elevates the professionalism of your PyTorch codebases significantly.

Key base classes in PyTorch you should be aware of:
```python
torch.Tensor
torch.optim.Optimizer
torch.utils.data.Dataset
torch.nn.Module
```

For more details, refer to the official PyTorch documentation or explore the PyTorch codebase using your code editor's symbol search feature.

`Dataclasses` are extremely useful for defining types. They allow for quick class definitions intended for data holding and validation, usually without an `__init__` method. Validation can be implemented using the `__post_init__` method.

Example:
```python
import torch
from dataclasses import dataclass

@dataclass
class PointCloud:
    points: torch.Tensor
    colors: torch.Tensor

    def __post_init__(self):
        assert self.points.shape == torch.Size([...])
        ...
```

Employing these types for network inputs and outputs significantly reduced my debugging time. I hope they will prove just as beneficial for you.

## Common Troubleshooting for DataLoaders

If you adopt a strict `dataclass` typing strategy in PyTorch, you might face a unique challenge. The PyTorch `DataLoader` typically expects data to consist of `Tensors` or iterable collections of `Tensors`. To resolve this, reimplement the default `collate_fn` and provide it to the `DataLoader`. Here's an example:

```python
@dataclass
class PointCloud:
    position: torch.Tensor
    color: torch.Tensor

def custom_collate_fn(batch: List[PointCloud]):
    positions = torch.stack([pcd.position for pcd in batch])
    colors = torch.stack([pcd.color for pcd in batch])
    return PointCloud(positions, colors)
...
loader = DataLoader(..., collate_fn=custom_collate_fn)
```

For distinguishing between batched and unbatched types, I usually rely on simple shape assertions.

With these minor adjustments, you'll find that your types work harmoniously with the PyTorch framework.

Let me know if you use types in your code and how they've benefited your projects.