PyTorch Linux训练脚本怎样编写

作者：袖梨 2026-07-04

编写一个在Linux环境下使用PyTorch进行训练的脚本涉及多个步骤。以下是一个简单的示例，展示了如何编写一个基本的训练脚本。

首先，确保你已经安装了PyTorch。你可以参考PyTorch官网的安装指南：https://pytorch.org/get-started/locally/

下面是一个简单的训练脚本示例：

import torchimport torch.nn as nnimport torch.optim as optimfrom torchvision import datasets, transformsfrom torch.utils.data import DataLoader# 定义一个简单的神经网络class SimpleNet(nn.Module):def __init__(self):super(SimpleNet, self).__init__()self.fc1 = nn.Linear(28 * 28, 128)self.fc2 = nn.Linear(128, 10)def forward(self, x):x = x.view(-1, 28 * 28)x = torch.relu(self.fc1(x))x = self.fc2(x)return x# 超参数batch_size = 64learning_rate = 0.001num_epochs = 5# 数据加载transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.1307,), (0.3081,))])train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)# 初始化模型、损失函数和优化器model = SimpleNet()criterion = nn.CrossEntropyLoss()optimizer = optim.Adam(model.parameters(), lr=learning_rate)# 训练模型for epoch in range(num_epochs):for i, (images, labels) in enumerate(train_loader):# 前向传播outputs = model(images)loss = criterion(outputs, labels)# 反向传播和优化optimizer.zero_grad()loss.backward()optimizer.step()if (i+1) % 100 == 0:print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{len(train_loader)}], Loss: {loss.item():.4f}')print('训练完成！')

解释

导入必要的库：
- torch 和 torch.nn 用于构建和训练神经网络。
- torch.optim 用于定义优化器。
- torchvision 用于加载和处理数据集。
- DataLoader 用于批量加载数据。
定义神经网络：
- SimpleNet 是一个简单的全连接神经网络，包含两个全连接层。
超参数：
- batch_size：每个批次的大小。
- learning_rate：学习率。
- num_epochs：训练的轮数。
数据加载：
- 使用 torchvision.transforms 对数据进行预处理。
- 加载 MNIST 数据集，并使用 DataLoader 进行批量加载。
初始化模型、损失函数和优化器：
- SimpleNet 实例化模型。
- CrossEntropyLoss 作为损失函数。
- Adam 作为优化器。
训练模型：
- 使用嵌套的 for 循环进行训练。
- 外层循环遍历每个 epoch。
- 内层循环遍历每个批次的数据。
- 进行前向传播、计算损失、反向传播和优化。