Transforms V2 proposal: Enabling reproducible workflows via local RNGs

[rsokl]

🚀 The feature

(This was originally pitched in this long feedback thread. It was recommended that I open a separate issue).

Enable the new transforms API to support the use of local generators to control RNG via the modified API:

from torch import Generator, Tensor, default_generator
import torch.nn as nn

class Transform(nn.Module):
    def _get_params(self, flat_inputs: List[Any], *, generator: Generator) -> Dict[str, Any]:
        ...

    def forward(self, *inputs: Any, generator: Generator = default_generator) -> Any:
        # the only modification
        params = self._get_params(flat_inputs, generator=generator)

Thus transforms that implement _get_params would replace calls like

# e.g. replace calls like
angle = float(torch.empty(1).uniform_(0.0, 180.).item())

with

# specifying the device is, unfortunately, necessary: https://github.com/pytorch/pytorch/issues/79018
angle = float(torch.empty(1, device=generator.device).uniform_(0.0, 180., generator=generator).item())

A transform like Compose would have to be modified as well. Currently, it supports a sequence of callables that are assumed to accept a single positional argument. It could be assumed that only instances of Transform involve stochasticity and will be passed the random generator. In this case, Compose would look like:

class Compose(Transform):
    # __init__ is unchanged

    def forward(self, *inputs: Any, generator: Generator = default_generator) -> Any:
        sample = inputs if len(inputs) > 1 else inputs[0]
        for transform in self.transforms:
            sample = transform(sample) if not isinstance(transform, Transform) else transform(sample, generator=generator)
        return sample

It would be straightforward to document this behavior to users – that only instances of Transform are passed the generator – so that they know how to opt-in to having the generator be passed to their custom transforms. And, again, this would be compatible with the old nn.Module transforms.

An example of this in practice would be:

from torch import Generator

rng = Generator.manual_seed(0)

trans = T.Compose(
    [
        T.ColorJitter(contrast=0.5),
        T.RandomRotation(30),
        T.CenterCrop(480),
    ]
)
img, bboxes, labels = trans(img, bboxes, labels, generator=rng)

Another nice thing about this is that specific fail cases that occur during training/testing can be reproduced in an isolated way; _get_params(dummy_img, generator=rng) can be used to iterate the generator's state to "replay" a sequence of transformations without have to redo all of the compute. Whereas this would not work if the model and the transforms both affect and derive from global state.

Motivation, pitch

In recent years, NumPy has completely revised their PRNG API to avoid global random state (here is a great post on good practices with NumPy's generators). JAX avoids mutable RNG objects altogether. PyTorch provides torch.Generator to users to to make randomness local and "non-spooky", but many libraries prevent users from utilizing this capability.

I am proposing that Transform enable users to optionally pass in a Generator to the forward pass so that torchvision transform pipelines can be made to be isolated from global entropy and thus support more reproducible workflows. This reproducibility is especially useful in the context of performing testing & evaluation – the specific sequence of data transformations performed should be able to be isolated from whether or not a model is using dropout in its forward pass.

Alternatives

No response

Additional context

@pmeier already provided (positive) feedback on this proposal here

cc @vfdev-5 @datumbox @bjuncek @pmeier

[vadimkantorov]

Also, an important thing is propagation of passed generators to all components of transform pipeline? (there's an additional complexity that not all transforms need these generators, but those that can accept them may need it to be propagated)

In my own code, I implemented my own containers and transform classes partly because of this. This is possible, but at least there should be reusable staticmethod ways for sampling the transform arguments that can accept rng.

[rsokl]

@vadimkantorov I updated my post to include a description of how one would pass rng through to all components of Compose-based pipeline.

[vadimkantorov]

@rsokl Yes, I implemented something similar in my own code. For that, all transforms must accept a generator argument even if they don't use any randomness and are deterministic, which may be an okay solution!

[pmeier]

Although somewhat niche and low priority, #3001 (comment) also shows an example of why good RNG support is needed. In case a user wants to use the same random parameters at different points in time, there are currently only two solutions:

1. Some random transformations in transforms v1 have the ability to sample the random parameters statically. However, there is no support to use these parameters directly and so also the parameters were known, the user had to write the actual transformation themselves. This might be a single call to a functional, but can be more complicated real quick.
2. The user could use the nuclear option and set torch.manual_seed(42) before each call. But without any containment this leaks into the surroundings.

Thus, in such a situation it would be really beneficial to just reset a generator and pass it again.

[vadimkantorov]

generalizing and accepting rng/generator optionally to those get_params / sample_params could be a first step towards easier reproducibility (if not done yet...)

[NicolasHug]

Passing the generator in forward() means that we can't have a per-transform RNG stream anymore when using containers like Compose() (and most use-cases involve Compose()). Taking the example from above

from torch import Generator

rng = Generator.manual_seed(0)

trans = T.Compose(
    [
        T.ColorJitter(contrast=0.5),
        T.RandomRotation(30),
        T.CenterCrop(480),
    ]
)
img, bboxes, labels = trans(img, bboxes, labels, generator=rng)

All 3 transforms will be using the same rng generator. There's no way to have a finer control where each transform would get its own RNG stream.

What if I want to "freeze" the RNG of one transform in Compose, while preserving maximal entropy for the rest of the transforms? I can't re-seed the rng object because that would affect all transforms. Is that something we want to enable?

In the limit, the only benefit I can see from allowing generators in forward() is that it separates the transforms RNG's from that of the global pytorch RNG (the one from torch.manual_seed()), but that's it. It still treats the transforms RNG as one (and only one) RNG. It doesn't allow a per-transform RNG in practice, since Compose must be used.

I'm curious if either @rsokl or @vadimkantorov have found this to be a limitation in practice?

(Instead of passing generators in forward(), we could pass them in __init__, but I need to think more about what that means.)

[NicolasHug]

Oh, I think we can rule-out forward() completely. How would that even work with datasets, where the tranforms forward is called in the dataset's __getitem__()?

[NicolasHug]

I spent a bit more time thinking about it and I implemented a toy solution to check what happens when DataLoader multi-processing is involved (passing the generators in __init__, as forward() is impossible). As far as I can tell, we're going to hit major UX issues (more details in the notebook below).

For now I can't think of a clean way to handle all this without requiring users to understand inner-details of the DataLoader, so that's a bummer. But I'm curious what you all think.

# %%
import torch
from torch.utils.data import DataLoader

class MyTransform(torch.nn.Module):
    def __init__(self, rng):
        super().__init__()
        self.rng = rng
    
    def forward(self):
        return torch.randint(0, 1000, size=(1,), generator=self.rng).item()

class Dataset:
    def __init__(self, transform):
        self.transform = transform

    def __getitem__(self, _):
        return self.transform()  # no input to the transform, we don't care.

    def __len__(self):
        return 1000

rng = torch.Generator()

t = MyTransform(rng)
ds = Dataset(t)


# %%
# Dataset only, so far so good
for x, _ in zip(ds, range(4)):
    print(x)
# 710
# 284
# 837
# 820

# %%
# Things break with DataLoder(num_workers > 0).
# The generator is duplicated across workers when we fork.
# Oopsies.
# Note: this is actually documented! https://pytorch.org/docs/stable/data.html#randomness-in-multi-process-data-loading
dl = DataLoader(ds, num_workers=2)
for x, _ in zip(dl, range(10)):
    print(x)
# tensor([299])
# tensor([299])
# tensor([754])
# tensor([754])
# tensor([334])
# tensor([334])
# tensor([739])
# tensor([739])
# tensor([609])
# tensor([609])
# %%

# Only way to make it work is to set a per-worker seed: https://pytorch.org/docs/stable/notes/faq.html#my-data-loader-workers-return-identical-random-numbers
# And BTW the only reason things "work" by default is becuase torch does that
# for us already  https://github.com/pytorch/pytorch/blob/1a661639f77a172df5d1ccd6987049292c6f3440/torch/utils/data/_utils/worker.py#L223-L225
def worker_init_fn(worker_id):
    rng.manual_seed(worker_id)

prev_state = rng.get_state()  # surpise surprise, see cell below
dl = DataLoader(ds, num_workers=2, worker_init_fn=worker_init_fn)
for x, _ in zip(dl, range(10)):
    print(x)
# tensor([44])
# tensor([845])
# tensor([239])
# tensor([139])
# tensor([933])
# tensor([124])
# tensor([760])
# tensor([368])
# tensor([963])
# tensor([263])


# %%
# Oh and on top of that, the RNG from the main process is never consumed!
# So we get the exact same RNG across epochs.
# EDIT: As Philip pointed out, this is actually already the case even when the global RNG is used everywhere
# Things only work OK because the global RNG gets consumed elsewhere e.g. by the RandomSampler. Ew.
assert (rng.get_state() == prev_state).all()