Mac Mini M2运行ESRGAN报错，如何调用M2 GPU进行深度学习？

报错原因 #

M2芯片属于Apple Silicon架构，仅支持Apple的Metal加速框架，而CUDA是NVIDIA GPU专属计算框架，因此代码默认检测CUDA会触发该断言错误。

解决步骤 #

• 安装适配Apple Silicon的PyTorch
  先卸载原有不支持MPS的PyTorch版本：

  pip uninstall torch torchvision torchaudio -y
  

  重新安装官方支持Metal的PyTorch：

  pip3 install torch torchvision torchaudio
  

• 修改ESRGAN代码切换计算设备
  找到代码中设置设备的片段，将CUDA检测逻辑替换为MPS检测：

  # 替换前
  # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  
  # 替换后
  device = torch.device('mps' if torch.backends.mps.is_available() else 'cpu')
  

  同时确保模型和输入数据都迁移到MPS设备上：

  model = model.to(device)
  input = input.to(device)
  

• 验证MPS是否正常工作
  在代码开头添加以下测试代码，确认MPS可用：

  import torch
  print("MPS可用:", torch.backends.mps.is_available())
  print("MPS已编译:", torch.backends.mps.is_built())
  

注意事项 #

• 确保Python环境是Apple Silicon原生版本（如Homebrew安装的Python、Conda的arm64版本），避免使用Rosetta模拟的x86环境
• 若代码中存在硬编码的torch.cuda.*相关操作，需替换为对应的torch.mps.*接口，PyTorch大部分通用张量操作会自动适配MPS

内容的提问来源于stack exchange，提问作者Cyan