MLP-Mixer基于自定义数据集训练时如何保存模型用于测试

一、训练完成后保存模型 #

这里提供两种常用的保存方案：

• 方案1：仅保存模型权重（推荐）
  该方案生成的文件体积小、兼容性强，不会受模型定义位置变动的影响，训练全部轮次完成后执行如下代码即可：

# 仅保存权重参数
torch.save(mixer_model.state_dict(), "mlp_mixer_weights.pth")

# 如果需要保存训练断点（方便后续恢复训练），可以保存为字典格式
torch.save({
    "epoch": epoch,
    "model_state_dict": mixer_model.state_dict(),
    "optimizer_state_dict": optimizer.state_dict(),
    "loss": current_loss,
}, "mlp_mixer_checkpoint.pth")

• 方案2：保存完整模型
  该方案无需在加载时重新定义模型结构，但要求加载环境中的MLPMixer类定义和保存时完全一致，否则会报错：

torch.save(mixer_model, "mlp_mixer_full.pth")

二、补全现有训练逻辑提示 #

你当前的代码还缺少损失函数、优化器、反向传播等核心训练逻辑，可参考如下代码补全后再执行保存操作：

if __name__ == '__main__':
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    dataset = customDataset()
    train_loader = DataLoader(dataset,batch_size=8,shuffle=True) # 可自行调整batch_size大小

    mixer_model = MLPMixer(in_channels=3, 
        image_size=416, 
        patch_size=16, 
        num_classes=2,
        dim=512, 
        depth=8, 
        token_dim=256, 
        channel_dim=2048).to(device)

    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(mixer_model.parameters(), lr=1e-4)
    num_epochs = 20 # 可根据训练效果调整训练轮次

    # 训练循环
    for epoch in range(num_epochs):
        mixer_model.train()
        running_loss = 0.0
        for i, data in enumerate(train_loader,0):
            inputs, labels = data
            inputs = inputs.to(device).float()
            labels = labels.to(device).squeeze() # 适配交叉熵损失的维度要求
            optimizer.zero_grad()
            outputs = mixer_model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
        print(f"Epoch {epoch+1}/{num_epochs}, Loss: {running_loss/len(train_loader):.4f}")
    
    # 训练完成后保存模型
    torch.save(mixer_model.state_dict(), "mlp_mixer_weights.pth")

三、加载模型用于测试 #

对应上面两种保存方案，加载方式分别如下：

对应方案1（仅存权重）的加载方法 #

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 先实例化和训练时参数完全一致的模型
mixer_model = MLPMixer(in_channels=3, 
    image_size=416, 
    patch_size=16, 
    num_classes=2,
    dim=512, 
    depth=8, 
    token_dim=256, 
    channel_dim=2048)
# 加载权重，GPU训练的模型要在CPU上加载时添加map_location参数即可
mixer_model.load_state_dict(torch.load("mlp_mixer_weights.pth", map_location=device))
mixer_model = mixer_model.to(device)
mixer_model.eval() # 切换为评估模式，关闭dropout等训练专属逻辑

对应方案2（存完整模型）的加载方法 #

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
mixer_model = torch.load("mlp_mixer_full.pth", map_location=device)
mixer_model.eval()

四、单张图像测试示例 #

import cv2
import torch

def predict_image(img_path, model, device):
    class_map = {0: "dogs", 1: "cats"}
    img_dim = (416, 416)
    # 预处理图像，和训练时的预处理逻辑保持一致
    img = cv2.imread(img_path)
    img = cv2.resize(img, img_dim)
    img_tensor = torch.from_numpy(img).permute(2, 0, 1).float()
    img_tensor = img_tensor.unsqueeze(0).to(device) # 添加batch维度
    
    with torch.no_grad(): # 关闭梯度计算，加快推理速度、节省显存
        output = model(img_tensor)
    pred_class_id = torch.argmax(output, dim=1).item()
    return class_map[pred_class_id]

# 调用示例
print(predict_image("test_img.jpg", mixer_model, device))

内容的提问来源于stack exchange，提问作者Simplicity