Pytorch-如何使用GPU进行训练

最新推荐文章于 2025-05-13 23:02:39 发布

wvdon

最新推荐文章于 2025-05-13 23:02:39 发布

阅读量9.4k

点赞数 2

CC 4.0 BY-SA版权

分类专栏： Pytorch 入门与实践文章标签： pytorch

本文链接：https://blog.youkuaiyun.com/WeDon_t/article/details/104300877

Pytorch 入门与实践专栏收录该内容

4 篇文章

订阅专栏

本文详细介绍如何在PyTorch中使用GPU进行高效计算，包括检查GPU状态、使用CUDA设备对象进行张量操作，以及通过实战案例展示如何利用GPU加速波士顿房价预测模型的训练过程。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

默认你已经安装了cuda，并且支持GPU。
如何查看你的GPU运行：

nvidia-smi

官方入门文档：使用gpu计算

# is_available 函数判断是否有cuda可以使用
# ``torch.device``将张量移动到指定的设备中
if torch.cuda.is_available():
    device = torch.device("cuda")          # a CUDA 设备对象
    y = torch.ones_like(x, device=device)  # 直接从GPU创建张量
    x = x.to(device)                       # 或者直接使用``.to("cuda")``将张量移动到cuda中
    z = x + y
    print(z)
    print(z.to("cpu", torch.double))

实战

波士顿预测，使用gpu进行训练：

加载数据

from sklearn.datasets import load_boston
boston = load_boston()
X,y   = (boston.data, boston.target)

切分数据

from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=0)

检测cuda是否可用

import torch 
from torch.utils import data
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
dtype = torch.FloatTensor
train_loader = data.DataLoader(myset,batch_size=128,shuffle=True)

建立网络

import torch.nn as nn
class Net1(nn.Module):
        def __init__(self, in_dim, n_hidden_1, n_hidden_2, out_dim):
            super(Net1, self).__init__()
            self.layer1 = torch.nn.Sequential(nn.Linear(in_dim, n_hidden_1))
            self.layer2 = torch.nn.Sequential(nn.Linear(n_hidden_1, n_hidden_2))
            self.layer3 = torch.nn.Sequential(nn.Linear(n_hidden_2, out_dim))
        def forward(self, x):
            x1 = F.relu(self.layer1(x))
            x1 = F.relu(self.layer2(x1))
            x2 = self.layer3(x1)
            print("\tIn Model: input size", x.size(),"output size", x2.size())
            return x2

实例化网络

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#实例化网络
model = Net1(13, 16, 32, 1)
if torch.cuda.device_count() > 1:
    print("Let's use", torch.cuda.device_count(), "GPUs")
    model = nn.DataParallel(model)
model.to(device)