pytorch安装cuda注意事项

最新推荐文章于 2026-01-07 17:22:05 发布

原创最新推荐文章于 2026-01-07 17:22:05 发布 · 283 阅读

3 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #人工智能 #python

台式机(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 time

import numpy as np
import torch


def 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 device


def 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版本的时间消耗的区别，