1 前言

  本项目将实现MobileNetV3的自定义数据集训练，以及推理模型部署。本篇内容采用猫狗数据集，2分类模型。
  C++端预测结果如下所示。

2 项目内容详细说明

3.1 训练及模型转换

  在python环境下进行模型的训练与格式转换。





  （1）train.py 实现模型训练；
  （2）detect.py 通过pt模型验证推理；
  （3）convert_to_onnx.py 将pt模型转换成onnx模型；
  （4）onnxpredict.py 通过onnx模型推理验证；

3.2 模型测试（C++）

  C++（Qt）环境下的onnx模型验证测试工程。

3 代码

3.1 train.py

  train.py实现如下所示。

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import torchvision.models as models
import os
import time
import copy
import matplotlib.pyplot as plt
from tqdm import tqdm

# 设置设备
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")

# 数据增强和归一化
data_transforms = {
    'train': transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(p=0.5),
        transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ]),
    'val': transforms.Compose([
        transforms.Resize((224,224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ]),
}

# 数据集路径
data_dir = '/media/ai/5c45cbac-396a-4328-b602-e47bc899eb89/ai/DX/ZheDang_MobileNetV3/datasets'
train_dir = os.path.join(data_dir, 'train')
val_dir = os.path.join(data_dir, 'val')

# 创建数据集
image_datasets = {
    'train': datasets.ImageFolder(train_dir, data_transforms['train']),
    'val': datasets.ImageFolder(val_dir, data_transforms['val'])
}

# 创建数据加载器
dataloaders = {
    'train': DataLoader(image_datasets['train'], batch_size=32, shuffle=True, num_workers=4),
    'val': DataLoader(image_datasets['val'], batch_size=32, shuffle=False, num_workers=4)
}

# 获取数据集信息
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
class_names = image_datasets['train'].classes
num_classes = len(class_names)

print(f"类别: {class_names}")
print(f"训练集大小: {dataset_sizes['train']}")
print(f"验证集大小: {dataset_sizes['val']}")


# 加载预训练的MobileNetV3模型
def create_model(num_classes=2):
    model = models.mobilenet_v3_small(pretrained=True)

    # 修改最后的分类层
    num_features = model.classifier[3].in_features
    model.classifier[3] = nn.Linear(num_features, num_classes)

    return model


# 创建模型
model = create_model(num_classes=num_classes)
model = model.to(device)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)


# 训练函数
def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
    since = time.time()

    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0

    # 记录训练过程
    train_loss_history = []
    val_loss_history = []
    train_acc_history = []
    val_acc_history = []

    for epoch in range(num_epochs):
        print(f'Epoch {epoch}/{num_epochs - 1}')
        print('-' * 10)

        # 每个epoch都有训练和验证阶段
        for phase in ['train', 'val']:
            if phase == 'train':
                model.train()  # 训练模式
            else:
                model.eval()  # 评估模式

            running_loss = 0.0
            running_corrects = 0

            # 使用tqdm显示进度条
            dataloader = dataloaders[phase]
            pbar = tqdm(dataloader, desc=f'{phase} Epoch {epoch}')

            # 迭代数据
            for inputs, labels in pbar:
                inputs = inputs.to(device)
                labels = labels.to(device)

                # 清零梯度
                optimizer.zero_grad()

                # 前向传播
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)

                    # 反向传播 + 优化（仅在训练阶段）
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                # 统计
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)

                # 更新进度条
                pbar.set_postfix({
                    'loss': f'{loss.item():.4f}',
                    'acc': f'{torch.sum(preds == labels.data).item() / inputs.size(0):.4f}'
                })

            if phase == 'train':
                scheduler.step()

            epoch_loss = running_loss / dataset_sizes[phase]
            epoch_acc = running_corrects.double() / dataset_sizes[phase]

            print(f'{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}')

            # 记录历史
            if phase == 'train':
                train_loss_history.append(epoch_loss)
                train_acc_history.append(epoch_acc.cpu().numpy())
            else:
                val_loss_history.append(epoch_loss)
                val_acc_history.append(epoch_acc.cpu().numpy())

            # 深度复制最佳模型
            if phase == 'val' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())

        print()

    time_elapsed = time.time() - since
    print(f'训练完成于 {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s')
    print(f'最佳验证准确率: {best_acc:.4f}')

    # 加载最佳模型权重
    model.load_state_dict(best_model_wts)

    return model, train_loss_history, val_loss_history, train_acc_history, val_acc_history


# 开始训练
print("开始训练模型...")
num_epochs = 100
model, train_loss, val_loss, train_acc, val_acc = train_model(
    model, criterion, optimizer, scheduler, num_epochs=num_epochs
)

# 保存模型
torch.save(model.state_dict(), 'mobileNetV3_best_model.pth')
print("模型已保存为 mobileNetV3_best_model.pth")

# 绘制训练曲线
plt.figure(figsize=(12, 5))

plt.subplot(1, 2, 1)
plt.plot(train_loss, label='Train Loss')
plt.plot(val_loss, label='Val Loss')
plt.title('Training and Validation Loss')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()

plt.subplot(1, 2, 2)
plt.plot(train_acc, label='Train Accuracy')
plt.plot(val_acc, label='Val Accuracy')
plt.title('Training and Validation Accuracy')
plt.xlabel('Epochs')
plt.ylabel('Accuracy')
plt.legend()

plt.tight_layout()
plt.savefig('training_curves.png')
plt.show()


# 测试函数
def test_model(model, dataloader):
    model.eval()
    correct = 0
    total = 0

    with torch.no_grad():
        for inputs, labels in dataloader['val']:
            inputs = inputs.to(device)
            labels = labels.to(device)

            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)

            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    accuracy = 100 * correct / total
    print(f'测试准确率: {accuracy:.2f}%')
    return accuracy


# 测试模型
print("测试模型性能...")
test_accuracy = test_model(model, dataloaders)
print(f"最终测试准确率: {test_accuracy:.2f}%")

3.2 detect.py

predict_image方法具体实现如下所示。

    def predict_image(self, image_path):
        """预测单张图像"""
        try:
            # 加载和预处理图像
            image = Image.open(image_path).convert('RGB')
            input_tensor = self.transform(image).unsqueeze(0).to(self.device)

            # 预测
            with torch.no_grad():
                outputs = self.model(input_tensor)
                probabilities = torch.softmax(outputs, dim=1)
                confidence, predicted = torch.max(probabilities, 1)

            # 获取结果
            class_name = self.class_names[predicted.item()]
            confidence_score = confidence.item()

            return class_name, confidence_score

        except Exception as e:
            print(f"预测时出错: {e}")
            return None, None

3.3 convert_to_onnx.py

  convert_to_onnx.py实现见第4章。

import torch
import torch.nn as nn
import torch.onnx
from torchvision import models
import onnx
import argparse
import os

class ONNXConverter:
    def __init__(self, model_path, num_classes=2):
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.model = self.load_model(model_path, num_classes)
        self.class_names = ['Cats', 'Dogs']

    def load_model(self, model_path, num_classes):
        """加载训练好的模型（与预测代码保持一致）"""
        model = models.mobilenet_v3_small(pretrained=False)

        # 修改最后的分类层，与训练时一致
        num_features = model.classifier[3].in_features
        model.classifier[3] = nn.Linear(num_features, num_classes)

        # 加载权重
        if os.path.exists(model_path):
            checkpoint = torch.load(model_path, map_location=self.device)
            
            # 处理不同的模型保存格式
            if isinstance(checkpoint, dict):
                if 'state_dict' in checkpoint:
                    state_dict = checkpoint['state_dict']
                elif 'model' in checkpoint:
                    state_dict = checkpoint['model']
                else:
                    state_dict = checkpoint
            else:
                state_dict = checkpoint
            
            # 处理多GPU训练保存的模型（移除'module.'前缀）
            if list(state_dict.keys())[0].startswith('module.'):
                from collections import OrderedDict
                new_state_dict = OrderedDict()
                for k, v in state_dict.items():
                    name = k[7:] if k.startswith('module.') else k
                    new_state_dict[name] = v
                state_dict = new_state_dict
            
            model.load_state_dict(state_dict)
            model = model.to(self.device)
            model.eval()  # 设置为评估模式
            print(f"模型已从 {model_path} 加载")
        else:
            raise FileNotFoundError(f"模型文件不存在: {model_path}")

        return model

    def convert_to_onnx(self, onnx_path, input_shape=(1, 3, 224, 224), opset_version=13):
        """
        将PyTorch模型转换为ONNX格式
        
        参数:
            onnx_path: 输出ONNX文件路径
            input_shape: 输入张量的形状 (batch, channels, height, width)
            opset_version: ONNX算子集版本
        """
        print(f"开始转换模型到ONNX格式...")
        print(f"输入形状: {input_shape}")
        print(f"ONNX算子集版本: {opset_version}")
        
        # 创建虚拟输入（与预测时的预处理一致）
        dummy_input = torch.randn(*input_shape, device=self.device)
        
        try:
            # 导出模型到ONNX
            torch.onnx.export(
                self.model,              # 要导出的模型
                dummy_input,            # 模型输入
                onnx_path,              # 输出文件路径
                export_params=True,     # 导出模型参数
                opset_version=opset_version,  # ONNX算子集版本
                do_constant_folding=True,  # 执行常量折叠优化
                input_names=['input'],   # 输入名称
                output_names=['output'], # 输出名称
                dynamic_axes={          # 动态轴配置
                    'input': {0: 'batch_size'}, 
                    'output': {0: 'batch_size'}
                },
                verbose=False
            )
            
            print(f"✓ ONNX模型已保存: {onnx_path}")
            
            # 验证ONNX模型
            self.validate_onnx_model(onnx_path, dummy_input)
            
            return True
            
        except Exception as e:
            print(f"❌ 转换失败: {e}")
            return False

    def validate_onnx_model(self, onnx_path, dummy_input):
        """验证ONNX模型的正确性"""
        try:
            # 加载并验证ONNX模型
            onnx_model = onnx.load(onnx_path)
            onnx.checker.check_model(onnx_model)
            print("✓ ONNX模型验证通过")
            
            # 验证推理一致性
            self.validate_inference_consistency(onnx_path, dummy_input)
            
        except Exception as e:
            print(f"❌ ONNX模型验证失败: {e}")

    def validate_inference_consistency(self, onnx_path, dummy_input):
        """验证PyTorch和ONNX推理结果的一致性"""
        try:
            import onnxruntime as ort
            import numpy as np
            
            # PyTorch推理
            with torch.no_grad():
                torch_output = self.model(dummy_input)
                torch_probs = torch.softmax(torch_output, dim=1)
                torch_conf, torch_pred = torch.max(torch_probs, 1)
            
            # ONNX Runtime推理
            ort_session = ort.InferenceSession(onnx_path)
            ort_inputs = {ort_session.get_inputs()[0].name: dummy_input.cpu().numpy()}
            ort_outputs = ort_session.run(None, ort_inputs)
            
            # 计算ONNX输出的概率和预测
            ort_output_tensor = torch.from_numpy(ort_outputs[0])
            ort_probs = torch.softmax(ort_output_tensor, dim=1)
            ort_conf, ort_pred = torch.max(ort_probs, 1)
            
            # 比较输出结果
            if torch.allclose(torch_output.cpu(), ort_output_tensor, rtol=1e-03, atol=1e-05):
                print("✓ PyTorch和ONNX Runtime输出一致")
                print(f"PyTorch预测: {self.class_names[torch_pred.item()]} (置信度: {torch_conf.item():.4f})")
                print(f"ONNX预测: {self.class_names[ort_pred.item()]} (置信度: {ort_conf.item():.4f})")
            else:
                print("⚠ 输出存在差异，但可能仍在可接受范围内")
                
        except ImportError:
            print("⚠ 未安装onnxruntime，跳过推理验证")
        except Exception as e:
            print(f"⚠ 推理验证失败: {e}")

def main():
    # 设置命令行参数
    parser = argparse.ArgumentParser(description='将PyTorch模型转换为ONNX格式')
    parser.add_argument('--model', '-m', default='mobileNetV3_best_model.pth',
                        help='输入的PyTorch模型文件路径 (默认: mobileNetV3_best_model.pth)')
    parser.add_argument('--output', '-o', default='model.onnx',
                        help='输出的ONNX文件路径 (默认: model.onnx)')
    parser.add_argument('--batch-size', '-b', type=int, default=1,
                        help='批量大小 (默认: 1)')
    parser.add_argument('--opset', type=int, default=13,
                        help='ONNX算子集版本 (默认: 13)')
    
    args = parser.parse_args()

    # 检查模型文件是否存在
    if not os.path.exists(args.model):
        print(f"错误: 模型文件不存在: {args.model}")
        return

    try:
        # 初始化转换器
        converter = ONNXConverter(args.model)
        
        # 设置输入形状（与预测代码的预处理一致）
        input_shape = (args.batch_size, 3, 224, 224)
        
        # 执行转换
        success = converter.convert_to_onnx(
            onnx_path=args.output,
            input_shape=input_shape,
            opset_version=args.opset
        )
        
        if success:
            print("\n🎉 转换完成！")
            print(f"输入: {args.model}")
            print(f"输出: {args.output}")
            print(f"输入形状: {input_shape}")
            print("\nONNX模型可以用于:")
            print("1. TensorRT加速推理")
            print("2. OpenVINO部署")
            print("3. TensorFlow Serving")
            print("4. 移动端部署")
        else:
            print("\n❌ 转换失败")
            
    except Exception as e:
        print(f"程序执行出错: {e}")

if __name__ == "__main__":
    main()

3.4 onnxpredict.py

  onnxpredict.py实现见第4章。

3.5 C++测试工程

  C++实现工程实现见第4章。

4 资源下载

  本案例中涉及到的所有代码请到此处下载https://download.csdn.net/download/wang_chao118/92681071。