Fira优化器快速入门指南：高效训练神经网络模型

2025-05-31 16:32:51作者：董斯意

前言

在深度学习领域，优化器的选择对模型训练效果有着至关重要的影响。Fira优化器作为一种新型的内存高效优化器，通过创新的参数压缩技术，能够在保持模型性能的同时显著降低内存占用。本文将详细介绍如何使用Fira优化器训练一个简单的全连接神经网络，并在MNIST数据集上进行测试。

环境准备

在开始之前，需要确保已安装必要的Python包：

pip install torch torchvision fira

这三个包分别提供了PyTorch深度学习框架、计算机视觉相关工具集以及Fira优化器本身。

模型定义

我们构建一个简单的全连接神经网络，包含一个隐藏层和一个输出层：

import torch.nn as nn
import torch.nn.functional as F

class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, 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 = F.relu(self.fc1(x))           # 隐藏层使用ReLU激活函数
        x = self.fc2(x)                   # 输出层不使用激活函数
        return x

这个网络结构简单但足以处理MNIST这样的基础分类任务。

数据准备

MNIST数据集包含60,000张训练图像和10,000张测试图像，每张都是28×28像素的手写数字：

from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# 定义数据转换：归一化处理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# 加载数据集
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, transform=transform)

# 创建数据加载器
train_loader = DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=64, shuffle=False)

模型初始化与优化器配置

Fira优化器的核心优势在于其内存高效性，通过divide_params函数可以灵活配置：

from fira import FiraAdamW, divide_params

model = SimpleNN()
criterion = nn.CrossEntropyLoss()  # 交叉熵损失函数

# 参数分组配置
param_groups = divide_params(
    model, 
    target_modules_list=["Linear"],  # 指定应用Fira的模块类型
    rank=8,                         # 压缩维度
    update_proj_gap=200,            # 梯度投影更新间隔
    alpha=1.0,                      # 学习率调整系数
    proj_type='std'                 # 投影类型
)

# 初始化Fira优化器
optimizer = FiraAdamW(
    param_groups, 
    lr=0.01,                       # 学习率
    betas=(0.9, 0.999),            # Adam的beta参数
    eps=1e-06,                     # 数值稳定性参数
    weight_decay=0.0,              # 权重衰减
    correct_bias=True              # 是否修正偏差
)

训练与测试流程

定义训练和测试函数，形成完整的训练循环：

def train(model, train_loader, criterion, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

        # 每100个batch打印一次训练进度
        if batch_idx % 100 == 0:
            print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)}] Loss: {loss.item():.6f}')

def test(model, test_loader, criterion):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            output = model(data)
            test_loss += criterion(output, target).item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)
    accuracy = 100. * correct / len(test_loader.dataset)
    print(f'\nTest set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({accuracy:.2f}%)\n')

# 训练5个epoch
for epoch in range(1, 6):
    train(model, train_loader, criterion, optimizer, epoch)
    test(model, test_loader, criterion)