pytorch安装

pytorch简介

• 具有强大GPU加速的张量计算（如NumPy）
• 基于基于磁带的自动刻度系统构建的深度神经网络
• torch (opens new window)

如果使用cpu推理直接pip install torch 即可，GPU推理继续阅读

安装

应该安装什么样的版本，首先需要参考使用的模型及显卡cuda版本，应该根据模型看兼容的cuda版本，根据cuda版本选择兼容的pytorch版本。

pytorch安装方法，按照系统选择适合的版本 (opens new window)

windows示例：

• 卸载当前版本pytorch，如果为避免不必要的麻烦，默认pip安装的torch只支持cpu

  最好是先安装pytorch这样可控，如果是先安装项目依赖可能会被动把pytorch下载下来，为防止版本不对就需要先进性卸载。当然也可以先执行最下方命令验证下是否支持gpu，如以支持则跳过安装。

  • pip uninstall torch torchvision torchaudio
    

• 安装11.8版本

  • pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
    

• 验证

  • pip list 有如下显示：

  • torch              2.0.1+cu118
    torchaudio         2.0.2+cu118
    torchvision        0.15.2+cu118
    

  • 进入python命令行

    • import torch
      >>> torch.cuda.is_available()
      True
      

    • 输出True则表示安装的torch支持GPU