📚 课程目录第3章 PyTorch模型基础

3. PyTorch模型基础

在Jetson平台上部署深度学习模型，核心链路是“PyTorch训练 → 模型导出 → TensorRT推理”。PyTorch作为动态图框架，训练灵活但推理效率不足；TensorRT作为静态图优化引擎，追求极致延迟与吞吐。两者之间的桥梁，就是本章要解决的三个核心问题：如何规范地定义模型、如何完成训练并导出为可移植格式、以及如何通过torch.jit将动态图转换为静态图。其中，torch.jit.trace与torch.jit.script的选择直接决定了导出成功率与推理精度，是工程落地中最容易踩坑的环节。

3.1 模型定义：从nn.Module到可导出结构

核心原理：动态图与静态图的本质差异

PyTorch的nn.Module采用动态图机制（Define-by-Run），每次前向传播都会重新构建计算图。这种设计带来了极大的灵活性——你可以在forward()中使用Python的if、for循环、甚至动态创建层。但代价是：推理时无法预知计算图结构，也就无法进行激进的算子融合、内存复用等静态优化。TensorRT恰恰需要完整的静态图描述。因此，模型定义阶段就要为后续导出做好准备。

  关键原则：在__init__中声明所有子模块和参数，在forward()中只描述数据流。避免在forward()中动态创建新的nn.Module或使用Python原生容器（如list/dict）存储模块——这些操作在trace时会被“冻结”为常量，在script时可能直接报错。

判断方法：模型可导出性自检清单

输入输出固定：确认输入张量的维度是否固定。若存在动态batch或动态分辨率，需在导出时明确指定动态轴。
无Python副作用：检查forward()中是否使用了print()、文件读写、随机数生成（非torch.randn）等操作。
算子兼容性：确认使用的所有算子（如F.interpolate、torch.cat）在TorchScript中均有对应实现。可提前用torch.jit.script做一次试转换。
设备无关性：模型参数应使用float32（或明确指定半精度），避免在forward()中硬编码.cuda()或.cpu()调用。

实战场景：一个典型分类模型的规范化定义

import torch
import torch.nn as nn

class SimpleCNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
        )
        self.classifier = nn.Linear(32 * 8 * 8, num_classes)

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)  # 使用view而非reshape，便于trace
        x = self.classifier(x)
        return x

这段代码遵循了“所有层在__init__中声明，forward()只做数据流操作”的原则。x.view(x.size(0), -1)中的x.size(0)在trace时会记录为动态维度，后续可映射为TensorRT的dynamic batch。

常见误区与总结

误区1：在forward()中使用len(x)或x.shape[0]等Python原生操作。这些在trace时会被记录为常量，导致导出模型只能处理固定batch。
误区2：使用nn.Dropout后忘记在导出前调用model.eval()。Dropout在训练模式下会随机丢弃神经元，trace时记录的是“丢弃后的路径”，导致推理结果错误。
误区3：将torch.Tensor作为nn.Module的属性直接赋值。正确做法是注册为buffer（self.register_buffer()）或Parameter。

3.2 训练与导出流程：从权重到可部署模型

核心原理：训练与导出的关注点分离

训练阶段追求的是模型精度，导出阶段追求的是推理效率。两者不应混为一谈。训练时使用model.train()启用Dropout和BatchNorm的统计更新；导出时必须切换到model.eval()，将BatchNorm的running_mean/running_var固定为训练期间累积的全局统计量。这一步骤的遗漏是导致“训练精度高、部署精度低”的首要原因。

操作流程：标准导出五步法

保存最佳权重：训练过程中监控验证集指标，保存model.state_dict()（仅权重）或完整模型（含结构）。推荐保存state_dict，便于后续加载到不同定义的模型中。
加载并切换模式：加载权重后，调用model.eval()。若模型包含BatchNorm，还需确认track_running_stats=True（默认值）。
构造示例输入：创建一个与实际推理输入shape一致的张量，如torch.randn(1, 3, 224, 224)。示例输入的数值不影响导出结果，但shape必须准确。
执行torch.jit转换：根据模型特性选择trace或script（详见3.3节）。
验证导出模型：用同一输入分别运行原始PyTorch模型和导出后的TorchScript模型，比较输出差异。通常允许1e-5级别的浮点误差。

案例拆解：一个包含动态控制流的模型导出

假设我们有一个目标检测模型，其forward()中包含根据输入尺寸动态调整anchor数量的逻辑：

class Detector(nn.Module):
    def forward(self, x):
        features = self.backbone(x)
        if x.shape[-1] > 100:  # 动态条件
            features = self.extra_layers(features)
        return self.head(features)

使用torch.jit.trace时，由于trace只记录实际执行路径，当示例输入尺寸为224×224时，if条件为True，extra_layers被记录；但部署时若输入变为128×128，条件为False，模型行为与训练时不一致。此时应改用torch.jit.script，它能解析if语句并生成真正的分支逻辑。

经验提醒：对于同时包含静态子网络和动态控制流的模型，推荐“混合方案”——先用torch.jit.script包装整个模型，再对内部静态子模块使用torch.jit.trace进行优化。这需要将子模块标记为@torch.jit.script_method或使用torch.jit.trace_module。

常见误区与总结

误区	后果	正确做法
导出前未调用`eval()`	BatchNorm使用batch统计量，推理精度下降	导出前显式调用`model.eval()`
示例输入shape与实际部署不一致	trace记录错误shape，部署时输入不匹配	使用与部署完全一致的输入shape
保存完整模型而非state_dict	模型结构耦合，难以迁移到不同环境	优先保存`state_dict`，加载时重建结构
忽略`torch.no_grad()`	导出时仍构建计算图，内存占用高	在导出代码块外包裹`with torch.no_grad():`

3.3 torch.jit.trace与torch.jit.script：机制、差异与选择

核心原理：两种转换范式的本质

torch.jit.trace采用“执行追踪”策略：给定一个示例输入，实际运行一次forward()，记录所有张量操作并固化为计算图。它不关心Python控制流，只记录“这次运行做了什么”。torch.jit.script则采用“源码解析”策略：直接解析forward()的Python源码，将其编译为TorchScript中间表示。它理解if、for、while等控制流，但要求代码风格符合TorchScript的语法子集。

  核心差异总结：trace是“数据驱动”的，适合结构固定、无数据依赖分支的模型；script是“代码驱动”的，适合包含动态控制流的模型。trace转换速度快、失败率低，但可能“静默地”忽略某些分支；script转换更严谨，但需要代码满足更多约束。

判断方法：五问决策法

问题1：模型forward()中是否有if/for/while依赖输入张量的值？是→script；否→trace。
问题2：模型是否使用了Python的list/dict推导式或zip/enumerate？是→优先script（trace可能记录错误）；否→trace。
问题3：模型是否包含nn.ModuleList或nn.ModuleDict？是→script更安全；否→trace。
问题4：是否需要在部署时支持动态输入shape？是→script（trace需要额外指定dynamic_axes）；否→trace。
问题5：模型是否包含自定义autograd.Function？是→两者都可能失败，需先实现torch.jit.script兼容版本；否→正常选择。

实战场景：trace与script的对比测试

import torch

# 场景A：静态网络（适合trace）
class StaticNet(nn.Module):
    def forward(self, x):
        return torch.relu(x + 1)

# 场景B：动态控制流（适合script）
class DynamicNet(nn.Module):
    def forward(self, x):
        if x.sum() > 0:
            return x * 2
        else:
            return x * 3

# trace测试
static_model = StaticNet()
traced = torch.jit.trace(static_model, torch.randn(4))
print("trace静态网络成功:", traced)

dynamic_model = DynamicNet()
try:
    traced_dynamic = torch.jit.trace(dynamic_model, torch.randn(4))
    print("trace动态网络成功（但可能有问题）:", traced_dynamic)
except Exception as e:
    print("trace动态网络失败:", e)

# script测试
scripted_dynamic = torch.jit.script(dynamic_model)
print("script动态网络成功:", scripted_dynamic)

运行结果会显示：trace对静态网络完全成功；对动态网络虽然能“成功”，但生成的图只包含x * 2分支（因为示例输入randn(4)的sum通常大于0）。而script能正确生成包含两个分支的图。

常见误区与总结

误区1：认为trace“失败”时才用script。实际上，trace可能“成功”但生成错误模型——它不会报错，只是静默地丢失分支。
误区2：script对Python语法支持有限。例如，yield、with语句、try/except均不支持。遇到此类代码，需重构为TorchScript兼容写法。
误区3：忽略torch.jit.trace的check_trace参数。设置check_trace=True（默认值）时，PyTorch会用额外输入验证trace结果，但仅检查数值一致性，不检查控制流完整性。

工程建议：在Jetson部署流程中，推荐“trace优先，script兜底”策略。先用trace快速验证模型能否导出，若发现控制流被错误冻结，再改用script。同时，无论使用哪种方式，导出后务必用多组不同输入（包括边界情况）验证输出一致性。

3.4 从TorchScript到TensorRT的衔接要点

完成torch.jit转换后，得到的.pt文件已包含静态计算图。但TensorRT无法直接消费TorchScript格式，通常需要先转换为ONNX，再通过trtexec或TensorRT Python API构建引擎。这一环节的常见问题包括：

算子映射缺失：某些PyTorch算子（如torch.einsum）在ONNX中没有直接对应，导出时可能报错或生成低效子图。建议在模型定义阶段就使用TensorRT支持的算子（如torch.matmul替代einsum）。
动态维度声明：在导出ONNX时，需通过dynamic_axes参数声明动态维度。例如：torch.onnx.export(model, dummy_input, "model.onnx", dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}})。
精度设置：TensorRT支持FP32、FP16和INT8。FP16通常能带来2-3倍加速，但需在导出ONNX时确保模型权重为FP32，由TensorRT在构建引擎时自动转换。

建议在Jetson上使用trtexec --onnx=model.onnx --saveEngine=model.trt --fp16命令行工具快速验证转换可行性，再集成到推理代码中。