📚 课程目录第2章深度学习模型基础回顾

深度学习模型基础回顾：PyTorch/TensorFlow模型定义、训练与导出、模型结构可视化

说实话，很多同学一上来就急着搞ONNX转TensorRT，结果模型导出来各种报错。我见过太多次了——模型结构写错了都不知道，导出时才发现算子不支持。所以这一章，咱们老老实实把基础打牢。

你想想看，模型转换的本质是什么？说白了就是把一种框架的“计算图”翻译成另一种框架能理解的格式。如果原始模型本身就写得稀里糊涂，那后面再怎么优化也是白搭。

2.1 PyTorch模型定义与训练

我个人习惯用PyTorch做原型开发，灵活、调试方便。咱们先看一个标准的模型定义流程。

2.1.1 模型定义

import torch
import torch.nn as nn

class SimpleCNN(nn.Module):
    def __init__(self, num_classes=10):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(16)
        self.relu = nn.ReLU(inplace=True)
        self.pool = nn.MaxPool2d(2)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(32)
        self.fc = nn.Linear(32 * 8 * 8, num_classes)

    def forward(self, x):
        x = self.pool(self.relu(self.bn1(self.conv1(x))))
        x = self.pool(self.relu(self.bn2(self.conv2(x))))
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

model = SimpleCNN()
print(model)

我的经验：定义模型时尽量把 inplace=True 用对地方。它节省显存，但导出ONNX时某些算子会报错。我建议导出前先关掉 inplace，省得后面排查半天。

2.1.2 训练与检查点保存

训练过程我就不啰嗦了，重点看保存方式。这里有个坑——很多人只保存 model.state_dict()，但做转换时你还需要模型结构。

# 推荐方式：保存完整模型
torch.save(model, 'model_full.pth')

# 或者保存 state_dict + 结构信息
torch.save({
    'model_state_dict': model.state_dict(),
    'input_shape': (1, 3, 32, 32),  # 记下输入尺寸
    'num_classes': 10
}, 'model_checkpoint.pth')

注意：ONNX导出时需要知道输入输出的名称和形状。我建议训练完第一时间用 torch.onnx.export 试导出，别等到部署时才发现问题。

2.2 TensorFlow模型定义与训练

TensorFlow这边，我个人更推荐用Keras API，简洁直观。不过要注意版本差异——TF1.x和TF2.x的导出方式完全不同。

2.2.1 模型定义（Keras方式）

import tensorflow as tf

model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(16, 3, padding='same', activation='relu', input_shape=(32, 32, 3)),
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.MaxPooling2D(2),
    tf.keras.layers.Conv2D(32, 3, padding='same', activation='relu'),
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.MaxPooling2D(2),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(10, activation='softmax')
])

model.summary()

2.2.2 模型保存与导出

TF2.x推荐用SavedModel格式，它包含了模型结构和权重，后续转ONNX最方便。

# 训练后保存
model.save('tf_model_savedmodel', save_format='tf')

# 或者保存为H5
model.save('tf_model.h5')

  关键点：TF转ONNX时，我遇到过很多自定义层不支持的情况。所以建议尽量用标准Keras层，少用Lambda层或自定义操作。如果非用不可，记得提前查一下ONNX算子支持列表。

2.3 模型导出为ONNX

嗯，这里才是重头戏。ONNX就像模型界的“通用语言”，但翻译过程中总会有信息丢失。

2.3.1 PyTorch导出ONNX

import torch.onnx

dummy_input = torch.randn(1, 3, 32, 32)
torch.onnx.export(
    model,
    dummy_input,
    'model.onnx',
    export_params=True,
    opset_version=11,  # 推荐11或更高
    do_constant_folding=True,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

避坑指南：我曾经在 opset_version 上吃过亏。版本太低有些算子不支持，版本太高TensorRT又不认识。我一般用opset 11或13，兼容性最好。

2.3.2 TensorFlow导出ONNX

TF这边需要用 tf2onnx 工具。安装很简单：pip install tf2onnx。

import tf2onnx

# 加载SavedModel
imported = tf.saved_model.load('tf_model_savedmodel')
spec = imported.signatures['serving_default']

# 转换
onnx_model, _ = tf2onnx.convert.from_function(
    spec,
    input_signature=[tf.TensorSpec([None, 32, 32, 3], tf.float32, name='input')],
    output_path='model.onnx'
)

2.4 模型结构可视化（Netron工具）

拿到ONNX文件后，第一件事就是看结构对不对。Netron是我最常用的工具，没有之一。

2.4.1 Netron使用方式

在线版：打开 netron.app，直接拖入模型文件
离线版：pip install netron，然后 netron model.onnx
支持格式：ONNX、PyTorch (.pth)、TensorFlow (.pb/.h5)、Keras、Caffe等

2.4.2 用Netron检查什么

检查项	说明	常见问题
输入输出节点	确认名称、形状、数据类型	输入名称不对，TensorRT绑定失败
算子类型	查看每个节点的op类型	出现不支持算子（如自定义OP）
数据流方向	跟踪张量流动路径	多余的分支或未使用的节点
常量折叠	检查是否有常量节点	常量过多影响TensorRT优化

我曾经踩过的坑：有一次模型导出后Netron显示一切正常，但TensorRT就是跑不起来。后来发现是某个Reshape节点的形状写死了，动态batch时直接崩了。所以一定要检查 dynamic_axes 是否生效。

2.5 本章知识体系

下面这张图总结了从模型定义到可视化的完整流程，你可以对照着检查自己的步骤。

2.6 本章小结

这一章我们走完了从模型定义到可视化的完整链路。你可能会觉得这些步骤很基础，但说实话，我见过太多人在导出ONNX这一步卡住——要么算子不支持，要么动态轴没配好，要么模型结构跟预期不一样。