使用 jupyter 在 PyTorch 上默认在 M1 GPU 上创建张量答案

【问题标题】：Creating tensors on M1 GPU by default on PyTorch using jupyter使用 jupyter 在 PyTorch 上默认在 M1 GPU 上创建张量
【发布时间】：2023-02-20 18:52:43
【问题描述】：

现在，如果我想在 gpu 上创建一个张量，我必须手动完成。对于上下文，我确信 GPU 支持是可用的，因为

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())

返回真。

我每次都这样做：

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

对于 jupyter notebook，有没有办法指定我所有的张量都应该在 GPU 上运行？

【问题讨论】：

标签： pytorch gpu

【解决方案1】：

方便的方式

根据 this issue 上的讨论，截至 2022 年 12 月 22 日，没有方便的方法将默认设备设置为 MPS。

不方便的方式

您可以通过拦截对 tensor constructors 的调用来实现“我不想为张量构造函数指定 device=，只需使用 MPS”的目标：

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

推荐方式：

我倾向于将您的设备放在笔记本顶部的配置中并明确使用它：

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

这需要您在整个代码中键入 device=Conf.dev。但是您可以轻松地将您的代码切换到不同的设备，并且您无需担心任何隐式全局状态。

【讨论】：

【解决方案2】：

从 2023-01-20 开始，每晚使用 pytorch 2.0，您可以使用以下方法将默认设备设置为 mps：

torch.set_default_device("mps")

【讨论】：