3. PyTorch基础回顾：张量操作、自动求导、nn.Module、模型保存与加载

好，咱们正式开始PyTorch的回顾。说实话，很多做FPGA部署的朋友，PyTorch基础其实挺扎实的，但一到模型导出、量化这些环节就容易翻车。我见过太多人，模型在GPU上跑得飞起，转成FINN就各种报错。问题出在哪？往往就是基础概念没吃透。

这一章，咱们把PyTorch最核心的四个东西过一遍。嗯，都是后面部署的基石。你想想看，如果连张量怎么变形、梯度怎么计算都不清楚，后面怎么去手动优化网络结构？

3.1 张量操作：一切的基础

张量，说白了就是PyTorch里的数组。但它比NumPy的ndarray多了两个杀手锏：GPU加速和自动求导。我个人习惯，拿到数据第一件事就是转成张量。

3.1.1 创建与基本属性

import torch

# 创建张量
x = torch.tensor([[1, 2], [3, 4]])
y = torch.zeros(3, 4)      # 全零
z = torch.ones(2, 3)       # 全一
w = torch.randn(2, 2)      # 标准正态分布

# 常用属性
print(x.shape)      # torch.Size([2, 2])
print(x.dtype)      # torch.int64
print(x.device)     # cpu

我在项目中遇到过一个问题：用torch.Tensor()创建张量，结果默认是float32，而我的模型需要int8。嗯，这里要注意，创建时最好显式指定dtype。

3.1.2 索引与切片

张量的索引和NumPy几乎一样。但有一个坑——张量切片返回的是视图，不是副本。我曾经因为这个，在修改切片数据时不小心改了原始张量，排查了半天。

# 基本索引
a = torch.randn(4, 3)
print(a[0])          # 第一行
print(a[:, 1])       # 第二列
print(a[1:3, :])     # 第二到第三行

# 高级索引（注意：返回的是副本）
rows = torch.tensor([0, 2])
cols = torch.tensor([1, 2])
print(a[rows][:, cols])

3.1.3 形状变换

形状变换是部署中最常用的操作。FINN对张量形状有严格要求，尤其是全连接层和卷积层的输入输出维度。

# 常用变换
x = torch.randn(2, 3, 4)
print(x.view(2, 12))          # 展平，要求内存连续
print(x.reshape(2, 12))       # 更灵活，自动处理不连续
print(x.permute(2, 0, 1))     # 维度重排，[4, 2, 3]
print(x.transpose(0, 1))      # 交换两维，[3, 2, 4]

# 增加/减少维度
print(x.unsqueeze(0).shape)   # [1, 2, 3, 4]
print(x.squeeze().shape)      # 移除所有大小为1的维度

你想想看，在FINN里，卷积层的输入是[N, C, H, W]，全连接层需要[N, C]。这中间的变换，全靠这些操作。

3.2 自动求导：训练的核心引擎

自动求导，说白了就是PyTorch帮你算梯度。你只需要定义前向计算，反向传播自动搞定。但部署时，我们更关心的是——如何关闭它？

3.2.1 计算图与梯度

# 开启梯度追踪
x = torch.ones(2, 2, requires_grad=True)
y = x + 2
z = y * y * 3
out = z.mean()

# 反向传播
out.backward()
print(x.grad)  # 查看梯度

我记得刚开始用PyTorch时，总搞不清requires_grad的作用。后来才明白，它决定了这个张量是否参与梯度计算。推理时一定要关掉，否则白白浪费内存和计算。

3.2.2 梯度控制

部署时，我们只需要前向推理。所以，关闭梯度是第一步。

# 方法1：全局关闭
torch.set_grad_enabled(False)

# 方法2：上下文管理器
with torch.no_grad():
    # 推理代码
    output = model(input)

# 方法3：detach() 分离张量
x = torch.randn(3, 4, requires_grad=True)
y = x.detach()  # y不再追踪梯度
print(y.requires_grad)  # False

3.3 nn.Module：构建网络的积木

nn.Module是所有网络层的基类。你写的每一个自定义层、每一个完整模型，都要继承它。说白了，它就是PyTorch的乐高积木。

3.3.1 基本用法

import torch.nn as nn
import torch.nn.functional as F

class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 10)
        self.dropout = nn.Dropout(0.5)

    def forward(self, x):
        x = x.view(x.size(0), -1)  # 展平
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x

这里有个细节：__init__里注册的层，PyTorch会自动管理参数。但如果你在forward里临时创建了nn.ReLU()，那它的参数就不会被注册。嗯，这会影响模型保存。

3.3.2 参数管理

部署时，我们经常需要查看或修改模型参数。比如，把权重转成定点数。

model = SimpleNet()

# 查看所有参数
for name, param in model.named_parameters():
    print(f"{name}: {param.shape}")

# 冻结参数（不更新梯度）
for param in model.parameters():
    param.requires_grad = False

# 获取特定层
fc1_weight = model.fc1.weight.data

3.4 模型保存与加载：部署的最后一公里

模型训练好了，怎么保存？怎么加载？看似简单，但部署时80%的问题都出在这里。我见过有人把整个模型对象序列化，结果换个环境就加载不了。

3.4.1 两种保存方式

方式	代码	适用场景
保存整个模型	torch.save(model, 'model.pth')	快速保存/加载，但依赖模型定义
保存状态字典	torch.save(model.state_dict(), 'model.pth')	推荐方式，跨环境兼容性好

# 推荐方式：保存state_dict
torch.save(model.state_dict(), 'simple_net.pth')

# 加载时，先创建模型实例，再加载参数
model = SimpleNet()
model.load_state_dict(torch.load('simple_net.pth'))
model.eval()  # 切换到推理模式

3.4.2 部署时的注意事项

模型加载后，一定要做两件事：model.eval()和关闭梯度。否则，Dropout和BatchNorm的行为会不一致。

# 完整的部署加载流程
def load_for_deployment(model_class, checkpoint_path):
    # 创建模型
    model = model_class()
    
    # 加载权重
    state_dict = torch.load(checkpoint_path, map_location='cpu')
    model.load_state_dict(state_dict)
    
    # 切换到推理模式
    model.eval()
    
    # 关闭梯度
    for param in model.parameters():
        param.requires_grad = False
    
    return model

你想想看，如果忘记eval()，Dropout层在推理时还会随机丢弃神经元，结果每次输出都不一样。这在部署中是绝对不能接受的。

3.5 本章小结

好，咱们快速回顾一下：

• 张量操作：创建、索引、形状变换，注意内存连续性和dtype
• 自动求导：理解计算图，部署时务必关闭梯度
• nn.Module：构建网络的基础，熟悉参数管理
• 模型保存与加载：推荐state_dict方式，记得eval()和关梯度

这些基础，后面在FINN的模型转换、量化、编译中会反复用到。我个人建议，如果你对某个点还不熟，花10分钟写个小demo跑一遍。嗯，动手永远是最好的学习方式。