使用GPU加速及配置

news2024/11/15 16:21:15

配置CUDA

英伟达

https://developer.nvidia.com/cuda-downloads

Python

python要求3.8.x版本以上

python下载

https://www.python.org/getit/

使用pytorch

查询地址：

https://pytorch.org/index.html

给出建议：

在这里插入图片描述

可以直接

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

shell中会报出版本信息，根据版本信息查找whl文件，手动下载

实在太慢，可以手动下载地址，建议使用迅雷下载

查询地址

https://download.pytorch.org/whl/torch/

安装

pip install path/文件名.whl

查看是否成功


import torch

print("PyTorch 版本:", torch.__version__)

if torch.cuda.is_available():
    print("CUDA 可用！")
    print("CUDA 版本:", torch.version.cuda)
    print("GPU 设备名称:", torch.cuda.get_device_name(0))
else:
    print("CUDA 不可用。")

# 创建一个张量
tensor = torch.tensor([1, 2, 3], dtype=torch.float32)

# 获取当前设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 将张量移动到 GPU 上
tensor_gpu = tensor.to(device)

print("张量已移动到设备:", tensor_gpu.device)
print("张量内容:", tensor_gpu)

# 创建两个张量，并将它们移动到GPU上
a = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32).to(device)
b = torch.tensor([[5, 6], [7, 8]], dtype=torch.float32).to(device)

# 矩阵乘法
matrix_product = torch.matmul(a, b)

# 将结果移回CPU打印
print("矩阵乘法的结果:")
print(matrix_product.cpu())