PyTorch Linux训练脚本如何编写

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

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

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

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from 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 = 64
learning_rate = 0.001
num_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。
- 内层循环遍历每个批次的数据。
- 进行前向传播、计算损失、反向传播和优化。

运行脚本

将上述代码保存为一个 Python 文件（例如 train.py），然后在终端中运行：

python train.py

确保你的环境中已经安装了所有必要的依赖项，并且数据集可以正确下载和加载。

解释

运行脚本

最新问答

相关标签