Pytorch如何指定device(cuda or cpu)例子解析

代码示例:

在PyTorch中,指定设备(CPU或CUDA)是一个非常重要的步骤,特别是当你在进行深度学习训练时。以下是一些指定设备的详细例子:

  1. 检查CUDA是否可用 :

    首先,你需要检查你的机器是否支持CUDA,并且PyTorch是否能够使用CUDA。

    python 复制代码
    import torch
    if torch.cuda.is_available():
        print("CUDA is available. Using GPU.")
    else:
        print("CUDA is not available. Using CPU.")
  2. 设置默认设备 :

    你可以设置PyTorch的默认设备,这样所有的张量和模型都会默认使用这个设备。

    python 复制代码
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  3. 创建张量并指定设备 :

    当你创建张量时,可以指定它们应该在哪个设备上。

    python 复制代码
    # 创建一个在CPU上的张量
    x_cpu = torch.tensor([1., 2., 3.], device='cpu')
    
    # 创建一个在GPU上的张量
    x_gpu = torch.tensor([1., 2., 3.], device=device)
  4. 将张量移动到指定设备 :

    如果张量已经创建,你可以使用.to().cuda()方法将其移动到指定的设备。

    python 复制代码
    # 将张量移动到GPU
    x_gpu = x_cpu.to(device)
    
    # 如果你知道你的设备是GPU,也可以使用.cuda()
    if torch.cuda.is_available():
        x_gpu = x_cpu.cuda()
  5. 指定模型的设备 :

    当你定义模型时,可以将其放置在指定的设备上。

    python 复制代码
    model = MyModel().to(device)
  6. 在训练循环中使用设备 :

    在训练循环中,你需要确保模型的输入数据和目标也在正确的设备上。

    python 复制代码
    for data, target in dataloader:
        data, target = data.to(device), target.to(device)
    
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
  7. 保存和加载模型时指定设备 :

    当你保存或加载模型时,确保模型在正确的设备上。

    python 复制代码
    # 保存模型
    torch.save(model.state_dict(), "model.pth")
    
    # 加载模型
    model = MyModel()
    model.load_state_dict(torch.load("model.pth", map_location=device))
    model.to(device)

请注意,当你在GPU上训练时,所有的输入数据、目标、模型参数等都应该在GPU上。这样可以确保计算是在GPU上进行的,从而提高训练速度。如果你的机器有多个GPU,你还可以指定使用特定的GPU,例如:

python 复制代码
device = torch.device("cuda:0")  # 使用第一个GPU

以上就是在PyTorch中指定设备的一些基本方法和例子。

喜欢本文,请点赞、收藏和关注!

相关推荐
Mrliu__5 分钟前
Python数据结构(七):Python 高级排序算法:希尔 快速 归并
数据结构·python·排序算法
C嘎嘎嵌入式开发10 分钟前
(22)100天python从入门到拿捏《【网络爬虫】网络基础与HTTP协议》
网络·爬虫·python
飞飞是甜咖啡11 分钟前
SPP-CNN解决CNN只能处理固定大小的输入图片
人工智能·神经网络·cnn
xiaoxiaode_shu15 分钟前
神经网络基础
人工智能·深度学习·神经网络
小小爱大王1 小时前
AI 编码效率提升 10 倍的秘密:Prompt 工程 + 工具链集成实战
java·javascript·人工智能
zzzyulin1 小时前
huggingface transformers调试问题--加载本地路径模型时pdb断点消失
python·transformer
教练、我想打篮球1 小时前
12 pyflink 的一个基础使用, 以及环境相关
python·flink·pyflink
飞翔的佩奇1 小时前
【完整源码+数据集+部署教程】【天线&运输】直升机战机类型识别目标检测系统源码&数据集全套:改进yolo11-CSP-EDLAN
前端·python·yolo·计算机视觉·数据集·yolo11·直升机战机类型识别目标检测系统
C嘎嘎嵌入式开发1 小时前
(21)100天python从入门到拿捏《XML 数据解析》
xml·开发语言·python
蓝博AI1 小时前
基于卷积神经网络的香蕉成熟度识别系统,resnet50,vgg16,resnet34【pytorch框架,python代码】
人工智能·pytorch·python·神经网络·cnn