PyTorch_指定运算设备 (包含安装 GPU 的 PyTorch)

PyTorch默认会将张量创建在 CPU 控制的内存中,即:默认的运算设备为 CPU。我们也可以将张量创建在 GPU 上,能够利用对于矩阵计算的优势加快模型训练。

将张量移动到 GPU 上有两种方法:

  1. 使用 cuda 方法
  2. 直接在 GPU 上创建张量
  3. 使用 to 方法指定设备

安装含有 GPU 的 PyTorch

通过这个可以判断电脑里是否已经安装 CUDA

python 复制代码
import torch 

print(torch.cuda.is_available()) # 判断是否有可用的 GPU 设备

如果结果输出是 False,说明设备里没有 GPU 的 PyTorch 的版本。

可以通过 PyTorch官网来安装含有 GPU 的 PyTorch 的版本。

安装 CUDA 12.8 版本的pip命令

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

这样就安装好含有 GPU 的 PyTorch 了。


代码

python 复制代码
import torch 

# 使用 cuda 方法
def test01():
    data = torch.tensor([10, 20, 30])
    print("存储设备:", data.device)

    # 将张量移动到 GPU 设备上
    data = data.cuda()
    print("存储设备:", data.device)
    
    # 将张量从 GPU 移动到 CPU 设备上
    data = data.cpu()
    print("存储设备:", data.device)

# 直接将张量创建在指定设备上
def test02():
    data = torch.tensor([10, 20, 30], device='cuda')
    print("存储设备:", data.device)

    # 将张量移动到 CPU 设备上
    data = data.cpu()
    print("存储设备:", data.device)

# 使用 to 方法
def test03():
    data = torch.tensor([10, 20, 30])
    print("存储设备:", data.device)

    # 将张量移动到 GPU 设备上
    data = data.to('cuda')
    print("存储设备:", data.device)
    
    # 将张量从 GPU 移动到 CPU 设备上
    data = data.to('cpu')
    print("存储设备:", data.device)


# 注意点:张量存储在不同设备上的张量不能够直接运算
def test04():
    data1 = torch.tensor([10,20,30])
    data2 = torch.tensor([1,2,3], device='cuda')

    # RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
    # data1 = data1.to('cuda') # 这样可以解决了
    data = data1 + data2 
    print(data)

    # 如果你的电脑上安装 pytorch 不是 gpu 版本的,或者电脑本身没有 gpu 设备环境
    # 否则下面的调用 cuda 函数的代码会报错
    # data1 = data1.cuda() 

if __name__ == "__main__":
    test04() 
相关推荐
麒羽7609 分钟前
PyTorch 实现 CIFAR10 数据集的 CNN 分类实践
pytorch·分类·cnn
热爱生活的猴子11 分钟前
使用bert或roberta模型做分类训练时,分类数据不平衡时,可以采取哪些优化的措施
人工智能·分类·bert
jie*15 分钟前
小杰机器学习高级(five)——分类算法的评估标准
人工智能·python·深度学习·神经网络·机器学习·分类·回归
nju_spy15 分钟前
计算机视觉 - 物体检测 开山鼻祖 R-CNN系列:Fast R-CNN、Faster R-CNN、Mask R-CNN
人工智能·计算机视觉·物体检测·fast r-cnn·r-cnn·mask r-cnn
视觉AI28 分钟前
如何实现 5 μm 精度的视觉检测?不仅仅是相机的事
人工智能·机器人·视觉检测·相机
这张生成的图像能检测吗32 分钟前
(论文速读)DiffBlender:可组合和通用的多模态文本到图像扩散模型
人工智能·深度学习·计算机视觉·文生图·扩散模型
清***鞋1 小时前
转行AI产品如何准备面试
人工智能·面试·职场和发展
高峰君主1 小时前
构建智能投资视野:用Python打造个性化股票分析系统
开发语言·python·股票
AI_Gump1 小时前
Qwen-image-edit在ComfyUI部署使用
人工智能·comfyui