一、操作流程
- 加载数据集
- 数据预处理:
- 将输入输出按特定格式拼接
- 文本转 Token IDs
- 通过 labels 标识出哪部分是输出(只有输出的 token 参与 loss 计算)
- 加载模型、Tokenizer
- 定义数据规整器
- 定义训练超参:学习率、批次大小、...
- 定义训练器
- 开始训练
- 注意:训练后推理时,输入数据的拼接方式要与训练时一致
二、训练样例
三、训练过程
1、 导入必要的库
from __future__ import print_function # 使print函数兼容Python 2.x
import torch # PyTorch的核心库
import torch.nn as nn # PyTorch的神经网络模块
import torch.nn.functional as F # PyTorch的函数模块,包含常用函数
import torch.optim as optim # PyTorch的优化器模块
from torchvision import datasets, transforms # 数据集和数据预处理工具
from torch.optim.lr_scheduler import StepLR # 学习率调度器
2、设置一些超参数
BATCH_SIZE = 64 # 每次训练的批量大小
TEST_BACTH_SIZE = 1000 # 每次测试的批量大小
EPOCHS = 5 # 训练的轮数(遍历整个训练集的次数)
LR = 0.01 # 学习率
SEED = 42 # 随机种子,确保结果可重复
LOG_INTERVAL = 100 # 每隔多少个批次打印一次日志
3、定义神经网络模型
class FeedForwardNet(nn.Module):
def __init__(self):
super().__init__()
# 第一层784维输入、256维输出 -- 图像大小28×28=784
self.fc1 = nn.Linear(784, 256)
# 第二层256维输入、128维输出
self.fc2 = nn.Linear(256, 128)
# 第三层128维输入、64维输出
self.fc3 = nn.Linear(128, 64)
# 第四层64维输入、10维输出 -- 输出类别10类(0,1,...9)
self.fc4 = nn.Linear(64, 10)
# Dropout module with 0.2 drop probability
self.dropout = nn.Dropout(p=0.2)
def forward(self, x):
# 把输入展平成1D向量
x = x.view(x.shape[0], -1)
# 每层激活函数是ReLU,额外加dropout
x = self.dropout(F.relu(self.fc1(x)))
x = self.dropout(F.relu(self.fc2(x)))
x = self.dropout(F.relu(self.fc3(x)))
# 输出为10维概率分布
x = F.log_softmax(self.fc4(x), dim=1)
return x
4、 定义训练过程
# 训练过程
def train(model, loss_fn, device, train_loader, optimizer, epoch):
# 开启梯度计算
model.train()
for batch_idx, (data_input, true_label) in enumerate(train_loader):
# 从数据加载器读取一个batch
# 把数据上载到GPU(如有)
data_input, true_label = data_input.to(device), true_label.to(device)
# 求解器初始化(每个batch初始化一次)
optimizer.zero_grad()
# 正向传播:模型由输入预测输出
output = model(data_input)
# 计算loss
loss = loss_fn(output, true_label)
# 反向传播:计算当前batch的loss的梯度
loss.backward()
# 由求解器根据梯度更新模型参数
optimizer.step()
# 间隔性的输出当前batch的训练loss
if batch_idx % LOG_INTERVAL == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data_input), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()
5、 定义测试过程
# 计算在测试集的准确率和loss
def test(model, loss_fn, device, test_loader):
model.eval() # 设置模型为评估模式
test_loss = 0 # 初始化测试损失
correct = 0 # 初始化正确预测数
with torch.no_grad(): # 关闭梯度计算
for data, target in test_loader: # 遍历测试数据集
data, target = data.to(device), target.to(device) # 将数据和标签移到设备
output = model(data) # 前向传播:计算输出
# sum up batch loss
test_loss += loss_fn(output, target, reduction='sum').item() # 累加批次损失
# get the index of the max log-probability
pred = output.argmax(dim=1, keepdim=True) # 获取最大概率的预测值
correct += pred.eq(target.view_as(pred)).sum().item() # 统计正确预测的数量
test_loss /= len(test_loader.dataset) # 计算平均测试损失
print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
test_loss, correct, len(test_loader.dataset),
100. * correct / len(test_loader.dataset))) # 打印测试结果:平均损失和准确率
6、 主函数
def main():
use_cuda = torch.cuda.is_available() # 检查是否有GPU可用
torch.manual_seed(SEED) # 设置随机种子,确保结果可重复
device = torch.device("cuda" if use_cuda else "cpu") # 设置计算设备:GPU或CPU
train_kwargs = {'batch_size': BATCH_SIZE} # 训练数据加载参数
test_kwargs = {'batch_size': TEST_BACTH_SIZE} # 测试数据加载参数
if use_cuda:
cuda_kwargs = {'num_workers': 1, 'pin_memory': True, 'shuffle': True} # 额外的CUDA参数
train_kwargs.update(cuda_kwargs) # 更新训练数据加载参数
test_kwargs.update(cuda_kwargs) # 更新测试数据加载参数
transform = transforms.Compose([ # 定义数据预处理步骤
transforms.ToTensor(), # 转为张量
transforms.Normalize((0.1307,), (0.3081,)) # 归一化
])
dataset_train = datasets.MNIST('data', train=True, download=True, transform=transform) # 下载并加载训练数据集
dataset_test = datasets.MNIST('data', train=False, transform=transform) # 下载并加载测试数据集
train_loader = torch.utils.data.DataLoader(dataset_train, **train_kwargs) # 定义训练数据加载器
test_loader = torch.utils.data.DataLoader(dataset_test, **test_kwargs) # 定义测试数据加载器
model = FeedForwardNet().to(device) # 创建神经网络模型,并移到设备
optimizer = optim.SGD(model.parameters(), lr=LR) # 定义优化器(随机梯度下降)
loss_fn = F.nll_loss # 定义损失函数(负对数似然损失)
for epoch in range(1, EPOCHS + 1): # 训练多个轮次
train(model, loss_fn, device, train_loader, optimizer, epoch) # 训练模型
test(model, loss_fn, device, test_loader) # 测试模型
7、 运行主函数
if __name__ == '__main__':
main() # 仅在脚本直接运行时执行主函数
四、 总结
整个代码通过以下几个步骤实现了一个完整的神经网络训练和测试流程:
1. **导入库**:引入必要的工具和库。
2. **设置参数**:定义训练和测试的相关超参数。
3. **定义模型**:构建一个简单的全连接神经网络模型。
4. **定义训练过程**:编写训练模型的逻辑。
5. **定义测试过程**:编写测试模型的逻辑。
6. **主函数**:设置设备、加载数据、创建模型、定义优化器和损失函数,然后进行训练和测试。
7. **运行主函数**:确保代码在直接运行时执行主函数。
通过这些步骤,我们实现了一个能够对MNIST手写数字进行分类的神经网络模型,并完成了训练和测试过程。