2、PyTorch模型基础：张量操作、模型定义、训练与保存（state_dict、完整模型）

在进入OpenVINO部署流程之前，必须扎实掌握PyTorch模型的生命周期管理。PyTorch作为动态图框架的代表，其核心操作围绕张量（Tensor）展开，模型定义遵循nn.Module的面向对象范式，训练过程依赖自动微分引擎，而模型保存则存在两种截然不同的策略——仅保存参数（state_dict）与保存完整模型。本章将系统拆解这四个环节，并重点剖析它们与后续OpenVINO转换之间的关联，因为部署阶段的诸多“坑”往往源于训练阶段对模型保存方式的理解偏差。

2.1 张量操作：PyTorch的“数据货币”

张量（Tensor）是PyTorch中最基础的数据结构，可以理解为支持GPU加速和自动微分的多维数组。在部署场景中，张量操作直接影响输入数据的预处理方式——OpenVINO要求输入为连续内存的NCHW或NHWC布局，而PyTorch默认的通道在前（NCHW）布局恰好与之兼容，但需要确保张量在permute或transpose后调用contiguous()，否则后续转换可能因内存不连续而报错。

核心原理：张量的四大属性

• shape（形状）：描述各维度大小，如[1, 3, 224, 224]表示1张3通道224×224图像。
• dtype（数据类型）：torch.float32、torch.int64等。OpenVINO对FP32和FP16支持最佳，INT8需额外校准。
• device（设备）：CPU或CUDA。模型保存时需注意设备一致性，否则加载时可能遇到“张量在GPU但模型在CPU”的错位。
• requires_grad（梯度需求）：训练时为True，推理时需通过torch.no_grad()或model.eval()关闭，以减少内存消耗。

操作流程：从图像到张量的标准流水线

import torch
from torchvision import transforms

# 预处理：PIL图像 → 张量 → 归一化 → 增加batch维度
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),          # HWC → CHW，像素值缩放到[0,1]
    transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                         std=[0.229, 0.224, 0.225])
])
input_tensor = transform(image)     # shape: [3, 224, 224]
input_batch = input_tensor.unsqueeze(0)  # shape: [1, 3, 224, 224]

# 关键：确保内存连续（OpenVINO转换前必须检查）
input_batch = input_batch.contiguous()
print(input_batch.shape, input_batch.dtype, input_batch.device)

2.2 模型定义：nn.Module的工程化范式

PyTorch模型通过继承torch.nn.Module类实现。所有层（卷积、全连接、激活函数）都在__init__中声明，前向计算逻辑写在forward方法中。这种设计使得模型结构清晰、参数自动管理，且天然支持嵌套组合。

判断方法：如何设计一个可部署的模型类

1. 层定义与forward分离：不要在forward中临时创建层，否则参数无法被model.parameters()捕获。
2. 使用nn.Sequential组织简单流水线：对于线性堆叠结构，Sequential可简化代码。
3. 动态控制流谨慎使用：if语句在forward中可用，但OpenVINO转换时可能产生额外子图，建议用torch.where或固定逻辑替代。

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(),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((1, 1))
        )
        self.classifier = nn.Linear(32, num_classes)
    
    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)  # 展平
        return self.classifier(x)

model = SimpleCNN()
print(model)
# 输出参数数量
total_params = sum(p.numel() for p in model.parameters())
print(f"总参数量: {total_params}")

案例拆解：模型定义中的常见陷阱

2.3 训练循环：从损失到收敛的完整闭环

训练过程本质上是优化问题：通过梯度下降最小化损失函数。PyTorch的自动微分引擎（Autograd）记录所有张量操作，在调用backward()时自动计算梯度，随后优化器更新参数。理解这一循环对部署至关重要——因为训练时的预处理、损失函数、优化器状态都不会被保存到部署文件中，只有模型参数被保留。

操作流程：标准训练循环的五个步骤

import torch.optim as optim

model = SimpleCNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(10):
    model.train()  # 启用训练模式（启用dropout等）
    running_loss = 0.0
    for images, labels in train_loader:
        # 1. 梯度清零
        optimizer.zero_grad()
        # 2. 前向传播
        outputs = model(images)
        loss = criterion(outputs, labels)
        # 3. 反向传播
        loss.backward()
        # 4. 参数更新
        optimizer.step()
        running_loss += loss.item()
    
    # 5. 验证评估
    model.eval()
    with torch.no_grad():
        correct = 0
        total = 0
        for images, labels in val_loader:
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    print(f"Epoch {epoch+1}, Loss: {running_loss/len(train_loader):.4f}, Acc: {100*correct/total:.2f}%")

2.4 模型保存：state_dict与完整模型的抉择

PyTorch提供两种持久化策略：仅保存参数（state_dict）和保存完整模型。两者在文件大小、加载方式、部署兼容性上差异显著。对于OpenVINO部署，强烈推荐使用state_dict方式，因为它只包含张量数据，不依赖Python类定义，便于跨环境迁移。

核心原理：两种保存方式的本质区别

操作流程：标准保存与加载范式

# 保存state_dict（推荐）
torch.save(model.state_dict(), 'model_weights.pth')

# 加载state_dict
model = SimpleCNN()  # 必须先实例化相同结构
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()

# 保存完整模型（不推荐用于部署）
torch.save(model, 'model_full.pth')
# 加载完整模型
model = torch.load('model_full.pth')  # 需要类定义可访问

案例拆解：一次部署中的模型加载失败复盘

背景与约束：某团队训练了一个ResNet18分类模型，在服务器上测试准确率92%。部署到边缘设备时，使用OpenVINO转换ONNX模型，但转换后推理结果完全错误。

初始现象：ONNX模型在OpenVINO Runtime中输出为随机概率分布，且置信度极低。

关键输入数据：训练时保存的是torch.save(model, 'resnet18_full.pth')，加载后直接用于导出。

分析判断过程：排查发现，完整模型保存时包含了训练模式下的BatchNorm状态。加载后未调用model.eval()，导致BN层使用batch统计量而非全局统计量。此外，完整模型加载依赖原始类定义，在部署环境中类路径不一致，导致加载时静默失败。

实际操作步骤：

1. 改用state_dict保存：torch.save(model.state_dict(), 'resnet18_weights.pth')
2. 加载时显式实例化模型：model = ResNet18(); model.load_state_dict(torch.load(...))
3. 调用model.eval()后再导出ONNX。

结果对比：转换后的ONNX模型在OpenVINO上准确率恢复至91.8%，推理速度提升2.3倍。

风险与代价：若使用完整模型保存，还需确保部署环境的Python版本、PyTorch版本与训练环境一致，否则可能触发反序列化错误。

复盘经验：部署链路中，模型保存方式是第一道关卡。始终使用state_dict，并在导出前验证model.training属性为False。