Adds functionality to populate torch generator using torch.thread_safe_generator

[divyanshk]

Added thread-safe random number generation to all V2 torchvision random transforms to prevent race conditions when using DataLoader with thread-based workers (worker_method='thread').

This is based on torch.thread_safe_generator which returns dataloader thread-worker specific RNG or None otherwise.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/vision/9371

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 1 Pending

As of commit a6bff65 with merge base 48956e0 ():

NEW FAILURE - The following job has failed:

• Lint / python-types / linux-job (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[NicolasHug]

Thanks for the PR @divyanshk . I think the changes look reasonable.

One thing I'm wondering is how does this affect the multiprocess-based dataloaders? Currently, since TV is using the global torch RNG, that global generator will be seeded by torch using a different seed for each process/worker. This is the correct behavior since we want each worker to have a different RNG.

Is that behavior preserved now that we're using torch.thread_safe_generator()?

It'd be good to have tests ensure that's the case (both for multiprocess and multithreaded cases).

[divyanshk]

The multiprocessing case remains unchanged because torch.thread_safe_generator will return None for multiprocessing use-case. So for MP, there is no change. Earlier the torch.rand functions received None for generator arg, and now they would get the same. Also added a test case where I confirm the expected behavior for multiprocessing.

[NicolasHug]

We have a few more random transforms in TV that we'll also want to update and test. I think the list you'll find in https://github.com/pytorch/vision/pull/7848/changes should have the proper coverage (but claude should be able to find all the relevant ones)

[divyanshk]

I was missing a lot of random transforms. Updated the PR to cover the transforms in these files in torchvision/transforms/v2:
_augment.py
_auto_augment.py
_color.py
_container.py
_geometry.py
_misc.py
_transform.py

This matches the files touched in your PR above.

RandomHorizontalFlip and RandomVerticalFlip were noisy for the test I want, i.e. two batches are different for different seeded workers. This is because the flipped outputs can be same for different seeds. So I added another reproducibility test which cover the flips and other transforms. This check if two transforms are the same for the same torch.thread_safe_generator value. This should add extra coverage.

[NicolasHug]

For this multi-threading test, is there a way to test the actual multi-threaded behavior, without the mocking? I.e. ideally I'd like to test the public-facing APIs when a user requests multi-threaded from the DataLoader. I'm not sure what the public entry point is though?

[divyanshk]

I agree. I am using the mocks only because the threading workers isn't landed yet and I want to land this PR first. Once in, I can update these tests. With respect to the transforms we are not simplifying anything though. With the mocking, each transform function is getting a different generator as would happen if done through the dataloading workers.