pytorch将数据与模型都放到GPU上训练

默认是CPU,如果想要用GPU需要:

  1. 安装配置cuda,然后更新/下载支持gpu版本的pytorch,可以参考:https://blog.csdn.net/weixin_35757704/article/details/124315569

  2. 设置device:

    py 复制代码
    device = torch.device('cuda' if torch.cuda.is_available else 'cpu')

    然后将数据与模型后面都额外加上.to(device)即可

示例程序

py 复制代码
import torch
import torch.nn as nn


# 一个简单的模型
class LinearRegressionModel(nn.Module):
    def __init__(self, input_shape, output_shape):
        super(LinearRegressionModel, self).__init__()
        self.linear = nn.Linear(input_shape, output_shape)

    def forward(self, x):
        out = self.linear(x)
        return out


def main():
    x_train = torch.randn(100, 4)  # 生成训练特征
    y_train = torch.randn(100, 1)  # 生成label
    model = LinearRegressionModel(x_train.shape[1], 1)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)  # 优化函数
    criterion = nn.MSELoss()  # 损失函数
    for epoch in range(100):
        optimizer.zero_grad()
        outputs = model(x_train)
        loss = criterion(outputs, y_train)
        loss.backward()
        optimizer.step()


if __name__ == '__main__':
    main()

修改为GPU版本:

py 复制代码
import torch
import torch.nn as nn


# 一个简单的模型
class LinearRegressionModel(nn.Module):
    def __init__(self, input_shape, output_shape):
        super(LinearRegressionModel, self).__init__()
        self.linear = nn.Linear(input_shape, output_shape)

    def forward(self, x):
        out = self.linear(x)
        return out


def main():
    # 1. 设置device
    device = torch.device('cuda' if torch.cuda.is_available else 'cpu')
    # 2. 数据与模型后都加 .to(device) 即可
    x_train = torch.randn(100, 4).to(device)  # 生成训练特征
    y_train = torch.randn(100, 1).to(device)  # 生成label
    model = LinearRegressionModel(x_train.shape[1], 1).to(device)  # next(transformer.parameters()).device

    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)  # 优化函数
    criterion = nn.MSELoss()  # 损失函数
    for epoch in range(100):
        optimizer.zero_grad()
        outputs = model(x_train)
        loss = criterion(outputs, y_train)
        loss.backward()
        optimizer.step()


if __name__ == '__main__':
    main()

修改后:

  1. 查看变量的位置:可以使用x_train.device查看tensor变量的位置
  2. 查看模型的位置:可以使用next(model.parameters()).device查看模型的位置

注意:不在同一个位置上的变量之间无法计算,模型无法使用不在同一个位置的数据

相关推荐
William.csj14 小时前
Pytorch/CUDA——flash-attn 库编译的 gcc 版本问题
pytorch·cuda
Green1Leaves1 天前
pytorch学习-9.多分类问题
人工智能·pytorch·学习
摸爬滚打李上进1 天前
重生学AI第十六集:线性层nn.Linear
人工智能·pytorch·python·神经网络·机器学习
HuashuiMu花水木1 天前
PyTorch笔记1----------Tensor(张量):基本概念、创建、属性、算数运算
人工智能·pytorch·笔记
喝过期的拉菲2 天前
如何使用 Pytorch Lightning 启用早停机制
pytorch·lightning·早停机制
kk爱闹2 天前
【挑战14天学完python和pytorch】- day01
android·pytorch·python
Yo_Becky2 天前
【PyTorch】PyTorch预训练模型缓存位置迁移,也可拓展应用于其他文件的迁移
人工智能·pytorch·经验分享·笔记·python·程序人生·其他
xinxiangwangzhi_2 天前
pytorch底层原理学习--PyTorch 架构梳理
人工智能·pytorch·架构
FF-Studio2 天前
【硬核数学 · LLM篇】3.1 Transformer之心:自注意力机制的线性代数解构《从零构建机器学习、深度学习到LLM的数学认知》
人工智能·pytorch·深度学习·线性代数·机器学习·数学建模·transformer
盼小辉丶3 天前
PyTorch实战(14)——条件生成对抗网络(conditional GAN,cGAN)
人工智能·pytorch·生成对抗网络