美文网首页程序员
如何在Mac上使用Metal加速PyTorch

如何在Mac上使用Metal加速PyTorch

作者: handyTOOL | 来源:发表于2022-09-18 16:35 被阅读0次

    概述

    之前在PC上,我们可以使用CUDA进行AI训练加速,但是在Mac上却只能使用CPU。现在苹果联合PyTorch,推出了Metal作为PyTorch的计算后端,苹果的文档描述如下

    Metal backend for PyTorch
    The new Metal backend in PyTorch version 1.12 enables high-performance, GPU-accelerated training using MPS Graph and the Metal Performance Shaders primitives.
    

    我们只需要在Mac上安装1.12版本以上的PyTorch,就可以使用MPS作为后端,享受Metal带来的计算加速了

    使用conda安装PyTorch

    基于conda安装torch和图像音频处理库

    conda install pytorch torchvision torchaudio -c pytorch
    

    测试MPS是否可用

    编写一个简单的例子测试MPS是否可用,在mps设备上创建一个tensor变量,然后打印出来

    import torch
    
    device = torch.device("mps")
    
    val = torch.rand((2,2), device=device)
    print(val)
    

    如果一切没问题,会打印出来类似的结果

    tensor([[0.9615, 0.5488],
            [0.1987, 0.4467]], device='mps:0')
    

    相关文章

      网友评论

        本文标题:如何在Mac上使用Metal加速PyTorch

        本文链接:https://www.haomeiwen.com/subject/xnfeortx.html