查看自己主机的显卡信息,比如我的

显卡是:NVIDIA GeForce GTX 1060 6GB

需要安装的工具列表:

  • 显卡驱动 (v465.89):
  • cuda v11.3.1
  • cuDNN v8.2.1
  • pytorch v1.10.1

步骤一:安装显卡驱动

使用下面的链接查找cuda版本对应的显卡驱动(此步骤千万不要弄错了)

https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/

NVIDIA 驱动程序下载 - 高级搜索

 

查看驱动版本相关信息 nvidia-smi

$nvidia-smi

 步骤二:安装cuda toolkit 11.3.1版本

官网:  

https://developer.nvidia.cn/zh-cn/cuda-toolkit

 

  • 设置环境变量NUMBA_ENABLE_CUDASIM

 步骤三:conda安装cudatoolkit(11.3.1)

$conda install -c cudatoolkit=11.3

步骤四:安装cuDNN

官网:https://developer.nvidia.cn/rdp/cudnn-archive

 

 点击下载,然后会得到一个压缩包,将压缩包解压出来的文件全部复制到之前装的CUDA路径下的里面就可。

 

步骤五:安装对应的Pytorch 

官网:

https://pytorch.org/get-started/locally/

$conda install pytorch torchvision cudatoolkit=11.3.1 -c pytorch

 成功时显示:

 查看pytorch版本

$python
>>> import torch
>>> print(torch.__version__)
 1.10.1

判断是否能使用cuda

$python
>>> import torch
>>> torch.cuda.is_available()
True

查看cuda版本

$nvcc --version

参考:

https://www.cnblogs.com/obarong/p/14833845.html

成功时训练打印:

Logo

欢迎来到FlagOS开发社区,这里是一个汇聚了AI开发者、数据科学家、机器学习爱好者以及业界专家的活力平台。我们致力于成为业内领先的Triton技术交流与应用分享的殿堂,为推动人工智能技术的普及与深化应用贡献力量。

更多推荐