Pytorch-07 完整训练测试过程

梦星辰.2024-05-29 18:50

要在PyTorch中使用GPU进行数据集的加载、模型的训练和最后模型的测试，需要将数据集和模型都移动到GPU上，并确保在训练和测试过程中都在GPU上进行计算。以下是一个完整的示例代码，展示了如何在PyTorch中使用GPU进行端到端的训练和测试：

python 复制代码

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# 检查GPU是否可用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 准备训练和测试数据，并将其移动到GPU
train_input = torch.randn(100, 10).to(device)
train_target = torch.randn(100, 1).to(device)
test_input = torch.randn(20, 10).to(device)
test_target = torch.randn(20, 1).to(device)

# 创建数据集和数据加载器
train_dataset = TensorDataset(train_input, train_target)
train_loader = DataLoader(train_dataset, batch_size=10, shuffle=True)

# 定义一个简单的神经网络模型，并将其移动到GPU
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc1 = nn.Linear(10, 5)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(5, 1)

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

model = SimpleModel().to(device)

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

# 训练模型
model.train()
for epoch in range(100):
    for input_data, target_data in train_loader:
        optimizer.zero_grad()
        output = model(input_data)
        loss = criterion(output, target_data)
        loss.backward()
        optimizer.step()

# 测试模型
model.eval()
with torch.no_grad():
    test_output = model(test_input)
    test_loss = criterion(test_output, test_target)
    print(f'Test Loss: {test_loss.item()}')

在这个示例中，我们首先检查GPU是否可用，并将训练和测试数据移动到GPU上。然后，我们创建了数据集和数据加载器，定义了神经网络模型，并将模型移动到GPU。在训练过程中，我们使用数据加载器加载数据进行训练；在测试过程中，我们使用model.eval()将模型切换为评估模式，并使用torch.no_grad()上下文管理器关闭梯度计算，以避免在测试过程中更新模型参数。最后，我们计算了模型在测试集上的损失。整个训练和测试过程都在GPU上进行，以加速计算和提高效率。