pytorch安装cuda注意事项

台式机(3070, 显存8G), 安装pytorch注意事项

pytorch 直接安装是没有cuda支持的,如果要用显卡(nvidia的cuda)支持，则要指定方式来处理

参考

https://docs.nvidia.com/cuda/cuda-installation-guide-linux/#post-installation-actions

1, 查看本地cuda版本,

如果是windows可以这时看:

或者 nvidia-smi (linux)

平时开发的过程中,要留意下显存占用

2, 下载nvidia的cuda工具包并安装

https://developer.nvidia.com/cuda-toolkit-archive

比如下载了12.8

参考文档安装即可

• 检查是否安装好了

打开命令行 nvcc -V

如果是linux 需要 先把路径加到环境里，如export PATH=${PATH}:/usr/local/cuda-12.8/bin

看nvidia-smi -L 情况
GPU 0: NVIDIA GeForce RTX 3070 (UUID: GPU-2fcc168a-79ff-7c0a-fca9-b5a927c89e0e)

这里是GPU 0 在一些部署如vllm之类的会指定哪个GPU来运行，一般默认是０, 如果有多卡可以留意一下.

3, pytorch 安装, 找对应的cuda版本来安装pytorch

这个不能直接安装，也要找对应的cuda版本来安装

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

比如本机:

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128

验证pytorch 中是否可用

torch.cuda.is_available()

测试性能:

import timeimport numpy as npimport torchdef get_cuda_device(): if torch.cuda.is_available(): torch.set_default_tensor_type(torch.cuda.FloatTensor) print(\"using cuda:\", torch.cuda.get_device_name(0)) pass device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\") return devicedef pyorch1(): a = torch.cuda.is_available() size = 10000 a = np.random.rand(size, size) b = np.random.rand(size, size) # begin = time.time() # x = np.dot(a, b) # print(f\"np:{time.time()-begin=}\") # begin = time.time() aa = torch.tensor(a) bb = torch.tensor(b) cc = torch.matmul(aa, bb) print(f\"cpu pytorch {time.time()-begin=}\") device = get_cuda_device() begin = time.time() aa = torch.tensor(a, device=device) bb = torch.tensor(b, device=device) cc = torch.matmul(aa, bb) print(f\"gpu {time.time()-begin=}\")if __name__ == \"__main__\": pyorch1()

上面能看到cpu版本和gpu版本的时间消耗的区别，