Pytorch-07 完整训练测试过程

要在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上进行,以加速计算和提高效率。

相关推荐
IT_陈寒2 小时前
React 18实战:7个被低估的Hooks技巧让你的开发效率提升50%
前端·人工智能·后端
逛逛GitHub3 小时前
飞书多维表“独立”了!功能强大的超出想象。
人工智能·github·产品
机器之心4 小时前
刚刚,DeepSeek-R1论文登上Nature封面,通讯作者梁文锋
人工智能·openai
CoovallyAIHub4 小时前
港大&字节重磅发布DanceGRPO:突破视觉生成RLHF瓶颈,多项任务性能提升超180%!
深度学习·算法·计算机视觉
CoovallyAIHub5 小时前
英伟达ViPE重磅发布!解决3D感知难题,SLAM+深度学习完美融合(附带数据集下载地址)
深度学习·算法·计算机视觉
aneasystone本尊6 小时前
学习 Chat2Graph 的知识库服务
人工智能
IT_陈寒6 小时前
Redis 性能翻倍的 7 个冷门技巧,第 5 个大多数人都不知道!
前端·人工智能·后端
飞哥数智坊17 小时前
GPT-5-Codex 发布,Codex 正在取代 Claude
人工智能·ai编程
倔强青铜三17 小时前
苦练Python第46天:文件写入与上下文管理器
人工智能·python·面试
虫无涯18 小时前
Dify Agent + AntV 实战:从 0 到 1 打造数据可视化解决方案
人工智能