pythonpytorchgpu

Creating tensors on M1 GPU by default on PyTorch using jupyter


Right now, if I want to create a tensor on gpu, I have to do it manually. For context, I'm sure that GPU support is available since

print(torch.backends.mps.is_available())# this ensures that the current current PyTorch installation was built with MPS activated.
print(torch.backends.mps.is_built())

returns True.

I've been doing this every time:

device = torch.device("mps")
a = torch.randn((), device=device, dtype=dtype)

Is there a way to specify, for a jupyter notebook, that all my tensors are supposed to be run on the GPU?


Solution

  • The convenient way

    There is no convenient way to set default device to MPS as of 2022-12-22, per discussion on this issue.


    The inconvenient way

    You can accomplish the objective of 'I don't want to specify device= for tensor constructors, just use MPS' by intercepting calls to tensor constructors:

    class MPSMode(torch.overrides.TorchFunctionMode):
        def __init__(self):
            # incomplete list; see link above for the full list
            self.constructors = {getattr(torch, x) for x in "empty ones arange eye full fill linspace rand randn randint randperm range zeros tensor as_tensor".split()}
        def __torch_function__(self, func, types, args=(), kwargs=None):
            if kwargs is None:
                kwargs = {}
            if func in self.constructors:
                if 'device' not in kwargs:
                    kwargs['device'] = 'mps'
            return func(*args, **kwargs)
    
    # sensible usage
    with MPSMode():
        print(torch.empty(1).device) # prints mps:0
    
    # sneaky usage
    MPSMode().__enter__()
    print(torch.empty(1).device) # prints mps:0
    

    The recommended way:

    I would lean towards just putting your device in a config at the top of your notebook and using it explicitly:

    class Conf: dev = torch.device("mps")
    # ...
    a = torch.randn(1, device=Conf.dev)
    

    This requires you to type device=Conf.dev throughout the code. But you can easily switch your code to different devices, and you don't have any implicit global state to worry about.