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.

❌ 2 New Failures, 1 Pending, 2 Unrelated Failures

As of commit 6277f11 with merge base 48956e0 ():

NEW FAILURES - The following jobs have failed:

• Build Windows Wheels / pytorch/vision / build-wheel-py3_10-cuda12_6 (gh)
  Process completed with exit code 1.
• Build Windows Wheels / pytorch/vision / build-wheel-py3_10-cuda12_8 (gh)
  Process completed with exit code 1.

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Build Windows Wheels / pytorch/vision / upload / upload-wheel-py3_10-cuda12_6 (gh) (trunk failure)
  Unable to download artifact(s): Artifact not found for name: pytorch_vision__3.10_cu126_x64
• Build Windows Wheels / pytorch/vision / upload / upload-wheel-py3_10-cuda12_8 (gh) (trunk failure)
  Unable to download artifact(s): Artifact not found for name: pytorch_vision__3.10_cu128_x64

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.