2. 基于知识蒸馏的VLA模型轻量化
任务说明
本案例面向机器人视觉-语言-动作（VLA）模型在边缘设备部署时的体积与延迟瓶颈，采用知识蒸馏（Knowledge Distillation）将大规模教师VLA模型的行为知识迁移至轻量学生网络。完成本案例后，你将获得一个可独立运行的蒸馏训练闭环，能够对比蒸馏前后学生模型的参数量与动作预测误差，验证压缩效果。
环境依赖
需要 Python 3.9+、PyTorch 2.0+、torchvision。若需GPU加速，请安装对应CUDA版本。无需额外数据集——本案例使用随机生成的低维模拟VLA特征与动作标签，聚焦蒸馏机制本身。
核心实现
以下代码定义了一个小型教师VLA网络（约1.2M参数）和一个学生网络（约0.3M参数），通过均方误差（MSE）蒸馏损失将教师的连续动作输出迁移给学生。训练循环中交替计算学生与教师输出的蒸馏损失，并报告验证误差。
```python import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader, TensorDataset # 固定随机种子保证可复现 torch.manual_seed(42) # 模拟VLA特征维度：视觉特征64维 + 语言指令嵌入32维 = 96维输入 INPUT_DIM = 96 # 动作空间：7自由度机械臂关节角度 + 1维夹爪开合 = 8维连续动作 ACTION_DIM = 8 BATCH_SIZE = 64 EPOCHS = 30 TEMPERATURE = 3.0 # 蒸馏温度系数 class TeacherVLA(nn.Module): """教师网络：3层MLP，参数量约1.2M""" def __init__(self): super().__init__() self.net = nn.Sequential( nn.Linear(INPUT_DIM, 512), nn.ReLU(), nn.Linear(512, 512), nn.ReLU(), nn.Linear(512, ACTION_DIM) ) def forward(self, x): return self.net(x) class StudentVLA(nn.Module): """学生网络：2层MLP，参数量约0.3M""" def __init__(self): super().__init__() self.net = nn.Sequential( nn.Linear(INPUT_DIM, 256), nn.ReLU(), nn.Linear(256, ACTION_DIM) ) def forward(self, x): return self.net(x) def generate_simulated_data(n_samples=2000): """生成模拟VLA特征与动作标签（确定性可复现）""" x = torch.randn(n_samples, INPUT_DIM) # 构造非线性映射作为“真实”动作生成规则 w = torch.linspace(-0.5, 0.5, ACTION_DIM).unsqueeze(1) y = torch.sin(x @ w.repeat(1, INPUT_DIM).T).sum(dim=1, keepdim=True) y = y.repeat(1, ACTION_DIM) + 0.1 * torch.randn(n_samples, ACTION_DIM) return x, y def train_student_with_distillation(): # 生成数据并划分训练/验证集 x, y = generate_simulated_data() split = int(0.8 * len(x)) train_loader = DataLoader(TensorDataset(x[:split], y[:split]), batch_size=BATCH_SIZE, shuffle=True) val_loader = DataLoader(TensorDataset(x[split:], y[split:]), batch_size=BATCH_SIZE) # 初始化教师与学生网络 teacher = TeacherVLA() student = StudentVLA() # 教师网络先在模拟数据上预训练（作为已收敛的强教师） teacher_opt = optim.Adam(teacher.parameters(), lr=1e-3) teacher_loss_fn = nn.MSELoss() for epoch in range(10): for xb, yb in train_loader: teacher_opt.zero_grad() loss = teacher_loss_fn(teacher(xb), yb) loss.backward() teacher_opt.step() teacher.eval() # 固定教师参数 # 蒸馏训练学生网络 student_opt = optim.Adam(student.parameters(), lr=1e-3) distill_loss_fn = nn.MSELoss() print("开始蒸馏训练...") for epoch in range(EPOCHS): student.train() total_loss = 0.0 for xb, yb in train_loader: student_opt.zero_grad() with torch.no_grad(): teacher_out = teacher(xb) student_out = student(xb) # 蒸馏损失：学生输出逼近教师输出（温度缩放后计算MSE） loss = distill_loss_fn(student_out / TEMPERATURE, teacher_out / TEMPERATURE) loss.backward() student_opt.step() total_loss += loss.item() * xb.size(0) avg_train_loss = total_loss / len(train_loader.dataset) # 验证：对比学生与教师输出差异 student.eval() val_diff = 0.0 with torch.no_grad(): for xb, yb in val_loader: s_out = student(xb) t_out = teacher(xb) val_diff += torch.mean((s_out - t_out) ** 2).item() * xb.size(0) avg_val_diff = val_diff / len(val_loader.dataset) if (epoch + 1) % 5 == 0: print(f"Epoch {epoch+1:2d} | 蒸馏损失: {avg_train_loss:.6f} | 学生-教师输出MSE: {avg_val_diff:.6f}") # 统计参数量 teacher_params = sum(p.numel() for p in teacher.parameters()) student_params = sum(p.numel() for p in student.parameters()) print(f"\n参数量对比: 教师 {teacher_params:,} → 学生 {student_params:,} | 压缩比 {teacher_params/student_params:.2f}x") return student, teacher if __name__ == "__main__": student, teacher = train_student_with_distillation() ```运行与验证
将上述代码保存为 vla_distill.py，在终端执行：
通过条件：
- 程序正常结束且无异常报错；
- 控制台输出参数量对比行，压缩比约为 4.0x（教师1.2M vs 学生0.3M）；
- 训练日志中每5轮的“学生-教师输出MSE”呈下降趋势，最终值低于0.05，表明学生已有效逼近教师输出。
关键说明
本案例使用模拟数据验证蒸馏流程的有效性，未涉及真实机器人数据集。实际应用中，将 generate_simulated_data() 替换为真实VLA样本（视觉编码器输出+语言嵌入拼接为96维向量，动作标签为8维连续值）即可直接迁移。温度系数 TEMPERATURE 控制软标签的平滑程度，值越大则学生更关注教师输出的类间相对关系；对于连续动作回归任务，MSE蒸馏损失已足够，无需使用KL散度。