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

相关推荐
慧星云16 分钟前
魔多 AI 支持 Kontext 在线训练 :超高角色一致性助您创作
人工智能·云计算·aigc
小周不长肉18 分钟前
AI幻觉的罪魁祸首
人工智能
小关会打代码21 分钟前
计算机视觉进阶教学之图像投影(透视)变换
人工智能·计算机视觉
星期天要睡觉22 分钟前
(纯新手教学)计算机视觉(opencv)实战十四——模板与多个对象匹配
人工智能·opencv·计算机视觉
一株月见草哇24 分钟前
(一)昇腾AI处理器技术
人工智能·ai·芯片
xieyan081125 分钟前
论文阅读_大模型情绪分析预测股票趋势
人工智能
mortimer31 分钟前
一次 ModelScope 替代 Hugging Face 的模型下载实战指南
人工智能·llm
柠檬味拥抱44 分钟前
2000 张人脸眼部检测数据集 | 标注规范、数据划分与应用场景
人工智能
芒果快进我嘴里1 小时前
超球损失函数
人工智能·机器学习