pytorch有GPU但是没调用

原创

mob64ca12e95b2b 2024-09-03 05:48:30 ©著作权

文章标签 CUDA 深度学习 Data 文章分类 PyTorch 人工智能

©著作权归作者所有：来自51CTO博客作者mob64ca12e95b2b的原创作品，请联系作者获取转载授权，否则将追究法律责任

PyTorch：如何确保利用GPU加速

在深度学习的实际应用中，计算速度是一个关键因素。PyTorch是一个非常流行的深度学习框架，支持GPU加速，但有时用户可能会遇到“PyTorch有GPU但是没调用”的情况。本文将分析这一问题，提供代码示例，并展示如何确保剪切计算时间。

一、为什么使用GPU？

GPU（图形处理单元）具有数量众多的核心和更高的并行处理能力，能够显著加速深度学习模型的训练。因此，利用GPU可以在相同时间内训练更多的模型或更复杂的模型，进而提升研究效率。

二、检查PyTorch是否能使用GPU

在开始之前，我们首先需要确认PyTorch是否检测到GPU。可以使用torch.cuda.is_available()来确认。

import torch

# 查看是否有可用的GPU
if torch.cuda.is_available():
    print("CUDA is available! Ready to use GPU.")
else:
    print("CUDA is not available! Using CPU instead.")

三、将模型和数据转移到GPU

一旦确认CUDA可用，就需确保将模型和数据转移到GPU。以下是一个简单的神经网络模型训练的示例，展示如何将其迁移到GPU。

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 定义简单的神经网络
class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(784, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        return self.fc2(x)

# 创建模型并转移到GPU
model = SimpleNN().to(device)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 加载数据
train_loader = torch.utils.data.DataLoader(
    datasets.MNIST('.', train=True, transform=transforms.ToTensor(), download=True),
    batch_size=64,
    shuffle=True
)

# 训练模型
for epoch in range(1, 3):  # 训练2个epoch
    model.train()
    for data, target in train_loader:
        data, target = data.view(-1, 784).to(device), target.to(device)
        
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        
    print(f'Epoch {epoch}, Loss: {loss.item()}')

在上述代码中，data和target都被转移到GPU，从而确保计算能够在GPU上进行。

四、甘特图示例

为了清晰展示GPU利用率的影响，以下是一个甘特图，用于展示训练过程中不同阶段的占用时间。

gantt
    title PyTorch GPU Utilization Steps
    dateFormat  YYYY-MM-DD
    section Data Preparation
    Load Data              :a1, 2023-10-20, 1d
    Data Preprocessing     :after a1  , 2d
    section Training
    Training Start         :a2, after a1 , 3d
    Model Evaluation       :after a2  , 1d

五、状态图示例

此外，我们可以使用状态图来展示PyTorch训练过程中的不同状态变化。

stateDiagram
    [*] --> Data_Preparation
    Data_Preparation --> Model_Training
    Model_Training --> Model_Evaluation
    Model_Evaluation --> [*]