[FR] torch context with default device & dtype

[ssnl]

🚀 Feature

There has been needs to set default device. This FR proposes an API that is compatible with 3rd party libs.

Motivation

There has been a lot of discussion around default device flag in pytorch (#7535). Yet, implementing such an API has been mostly blocked by the concern that 3rd party libraries may assume that tensors are created on CPU.

Similar dilemma has been also seen in the python multiprocessing library, where multiple start methods can be used (triggering bugs like librosa/librosa#747). They come up with this API multiprocessing.get_context(xxx), which returns a context object with the same set of functions as the multiprocessing module, but associated with a different start method, enabling patterns like mp = mp.get_context('spawn').

In addition to tensor creation, for many users (including me) it is extremely verbose to write .to(device) for every tensor yielded tensor from the data loader. It would be very handy if this handles the moving automatically as well.

Pitch

1. Tensor creation

   import torch
   torch.empty(3)  # returns CPU float tensor
   torch = torch.get_context(device='cuda:2', dtype=torch.half)
   torch.empty(3)  # returns CUDA half tensor

2. Module creation

   import torch
   torch = torch.get_context(device='cuda:2')
   MyModule()  # returns module on CUDA if MyModule uses torch tensor creation methods

3. (needs discussion) affect torch_ctx.utils.data.DataLoader s.t. the yielded samples contain tensors moved to the device of torch_ctx.
4. (needs discussion) affect torch_ctx.load(xxx) the same way as map_location.

Concerns

Such features always have the potential issue of making things too "frictionless", obscure, and harder to debug. But I think this one is not too bad.

cc @ezyang @gchanan @zou3519 @bdhirsh @heitorschueroff @ngimel @ejguan

[soumith]

after one framework finally got rid of a global context after 3 years, are you saying the competing framework should add one? :P

[ssnl]

:D a different kind of context

[gchanan]

this sounds like a promising API -- there's probably some implementation complexity behind it, but we'll remove "needs research" as it looks acceptable.

[elistevens]

Unclear if this issue is still actually active, but it was pointed at recently as if it were, and I'd like to point out that #2 above would require some magic to make work properly (assuming that the torch = part is required, vs. just setting a bunch of global state in torch.get_context(device='cuda:2'), and in that case, the function should be renamed to something more accurate), and would leave a bunch of polluted global state behind, possibly just from importing a 3rd party package that used it.

Please use a context manager like with torch.set_default_device('cuda:0'):.

[ssnl]

@elistevens Yeah I think you are right. #2 doesn't work normally, unfortunately.