📚 课程目录第3章实现视觉编码器多任务共享

3. 实现视觉编码器多任务共享

任务目标

在多任务VLA（Vision-Language-Action）模型中，不同任务（如抓取、推拉、放置）共享同一视觉输入空间，但任务语义差异大。本案例实现一个基于预训练ViT的共享视觉编码器，使多个任务头复用同一套视觉特征提取参数，并验证特征是否具备任务通用性。

实现方案

采用冻结ViT-B/16骨干 + 可训练轻量适配层的共享编码器结构。ViT负责通用视觉特征提取，适配层（1×1卷积 + LayerNorm）将特征映射到任务无关的共享空间。三个任务头（抓取、推拉、放置）分别从共享特征中解码各自的动作参数。

核心代码

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import ViTModel

class SharedVisionEncoder(nn.Module):
    """多任务共享视觉编码器：ViT骨干 + 轻量适配层"""
    def __init__(self, vit_model_name='google/vit-base-patch16-224-in21k',
                 feat_dim=256, freeze_backbone=True):
        super().__init__()
        # 加载预训练ViT，仅使用encoder输出
        self.vit = ViTModel.from_pretrained(vit_model_name)
        if freeze_backbone:
            for p in self.vit.parameters():
                p.requires_grad = False

        # 共享适配层：将768维patch特征压缩到256维共享空间
        self.adaptor = nn.Sequential(
            nn.Conv1d(768, feat_dim, kernel_size=1),
            nn.LayerNorm(feat_dim),
            nn.GELU()
        )

    def forward(self, pixel_values):
        # pixel_values: [B, 3, 224, 224]
        outputs = self.vit(pixel_values, output_hidden_states=False)
        # last_hidden_state: [B, 197, 768] (含[CLS])
        patch_tokens = outputs.last_hidden_state  # [B, 197, 768]
        # 转置适配: [B, 768, 197] → conv1d → [B, 256, 197]
        shared_feat = self.adaptor(patch_tokens.transpose(1, 2))
        # 转回 [B, 197, 256]
        return shared_feat.transpose(1, 2)

class TaskHead(nn.Module):
    """轻量任务头：从共享特征解码任务特定输出"""
    def __init__(self, feat_dim=256, hidden_dim=128, out_dim=7):
        super().__init__()
        self.fc = nn.Sequential(
            nn.Linear(feat_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, out_dim)
        )

    def forward(self, shared_feat):
        # 取[CLS] 对应的特征
        cls_feat = shared_feat[:, 0, :]  # [B, 256]
        return self.fc(cls_feat)

class MultiTaskVLA(nn.Module):
    """多任务VLA模型：共享编码器 + 三个任务头"""
    def __init__(self):
        super().__init__()
        self.encoder = SharedVisionEncoder()
        self.grasp_head = TaskHead(out_dim=7)    # 抓取: 位置(3)+旋转(4)
        self.push_head = TaskHead(out_dim=4)     # 推拉: 方向(2)+位移(2)
        self.place_head = TaskHead(out_dim=3)    # 放置: 位置(3)

    def forward(self, pixel_values, task_type):
        shared_feat = self.encoder(pixel_values)
        if task_type == 'grasp':
            return self.grasp_head(shared_feat)
        elif task_type == 'push':
            return self.push_head(shared_feat)
        elif task_type == 'place':
            return self.place_head(shared_feat)
        else:
            raise ValueError(f"Unknown task: {task_type}")

# ========== 最小验证闭环 ==========
if __name__ == '__main__':
    # 依赖: pip install torch transformers
    model = MultiTaskVLA()
    # 模拟一个batch的输入图像 (随机噪声即可验证前向传播)
    dummy_img = torch.randn(2, 3, 224, 224)

    # 验证三个任务头共享同一编码器输出
    with torch.no_grad():
        feat = model.encoder(dummy_img)
        print(f"共享特征形状: {feat.shape}")  # 预期: [2, 197, 256]

        grasp_out = model(dummy_img, 'grasp')
        push_out = model(dummy_img, 'push')
        place_out = model(dummy_img, 'place')

        print(f"抓取输出: {grasp_out.shape}")  # 预期: [2, 7]
        print(f"推拉输出: {push_out.shape}")   # 预期: [2, 4]
        print(f"放置输出: {place_out.shape}")  # 预期: [2, 3]

    # 验证共享性：三个任务头的梯度都流向同一适配层
    grasp_out.mean().backward()
    adaptor_grad = model.encoder.adaptor[0].weight.grad
    assert adaptor_grad is not None, "适配层应收到梯度"
    print("✓ 共享编码器参数成功接收梯度，多任务共享机制有效")

运行与验证

环境依赖：pip install torch transformers
运行方式：直接执行脚本 python shared_encoder.py
通过条件：
- 共享特征输出形状为 [2, 197, 256]，三个任务头输出维度分别为7、4、3；
- 反向传播后适配层梯度非空，证明三个任务头确实通过共享层回传梯度；
- ViT骨干参数 requires_grad=False，训练时仅更新适配层与任务头。

关键说明

ViT输出197个（1个[CLS] + 196个patch），本实现保留全部供后续空间操作使用；任务头仅取[CLS]做全局决策。
冻结ViT可大幅降低显存占用与训练时间，适配层仅约0.2M可训练参数，适合多任务联合训练场景。
实际训练时需将三个任务的数据按batch混合，通过task_type参数路由到对应任务头，实现参数共享。