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

最新推荐文章于 2025-07-14 11:24:10 发布

呆萌的代Ma

最新推荐文章于 2025-07-14 11:24:10 发布

阅读量856

点赞数 5

CC 4.0 BY-SA版权

分类专栏： pytorch/神经网络文章标签： pytorch

本文为优快云博主"呆萌的代Ma"原创文章，转载请注明博客链接：https://blog.youkuaiyun.com/weixin_35757704/

本文链接：https://blog.youkuaiyun.com/weixin_35757704/article/details/144770300

默认是CPU，如果想要用GPU需要：

安装配置cuda，然后更新/下载支持gpu版本的pytorch，可以参考：https://blog.youkuaiyun.com/weixin_35757704/article/details/124315569
设置device：
```
device = torch.device('cuda' if torch.cuda.is_available else 'cpu')
```
然后将数据与模型后面都额外加上.to(device)即可

示例程序

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