SmolVLA 模型微调

概述

SmolVLA（Small Vision-Language-Action）是由 HuggingFace 开发的轻量级视觉-语言-动作模型，专为机器人学习任务设计。该模型仅有 450M 参数，适合在消费级硬件上运行，是进行机器人学习研究和开发的理想选择。

先决条件

系统要求

• 操作系统：Linux（推荐 Ubuntu 20.04+）或 macOS
• Python 版本：3.8+
• GPU：NVIDIA GPU（推荐 RTX 3080 或更高），至少 8GB 显存
• 内存：至少 16GB RAM
• 存储空间：至少 50GB 可用空间

环境准备

1. 安装 LeRobot

# 克隆 LeRobot 仓库
git clone https://github.com/huggingface/lerobot.git
cd lerobot

# 创建虚拟环境
conda create -n lerobot python=3.10
conda activate lerobot

# 安装依赖
pip install -e .

2. 安装额外依赖

# 安装 Flash Attention（可选，用于加速训练）
pip install flash-attn --no-build-isolation

# 安装 Weights & Biases（用于实验跟踪）
pip install wandb
wandb login

数据准备

LeRobot 格式数据

SmolVLA 需要使用 LeRobot 格式的数据集。确保你的数据集包含以下结构：

your_dataset/
├── data/
│   ├── chunk-001/
│   │   ├── observation.images.cam_high.png
│   │   ├── observation.images.cam_low.png
│   │   └── ...
│   └── chunk-002/
│       └── ...
├── meta.json
├── stats.safetensors
└── videos/
    ├── episode_000000.mp4
    └── ...

数据质量要求

根据 HuggingFace 的建议，SmolVLA 需要：

• 最少 25 个高质量 episode 才能获得良好性能
• 推荐 100+ episode 以获得最佳效果
• 每个 episode 应包含完整的任务执行过程
• 图像分辨率推荐 224x224 或 256x256

微调训练

基本训练命令

# 设置环境变量
export HF_USER="io-ai-data"
export CUDA_VISIBLE_DEVICES=0

# 启动 SmolVLA 微调
lerobot-train \
  --policy.path=lerobot/smolvla_base \
  --dataset.repo_id=${HF_USER}/my_dataset \
  --dataset.root=/data/lerobot_dataset \
  --batch_size=64 \
  --steps=20000 \
  --output_dir=outputs/train/smolvla_finetuned \
  --job_name=smolvla_finetuning \
  --policy.device=cuda \
  --training.learning_rate=1e-4 \
  --training.warmup_steps=1000 \
  --wandb.enable=true \
  --wandb.project=smolvla_finetuning

高级训练配置

多 GPU 训练

# 使用 torchrun 进行多 GPU 训练
torchrun --nproc_per_node=2 --master_port=29500 \
  $(which lerobot-train) \
  --policy.path=lerobot/smolvla_base \
  --dataset.repo_id=${HF_USER}/my_dataset \
  --dataset.root=/data/my_dataset \
  --batch_size=32 \
  --steps=20000 \
  --output_dir=outputs/train/smolvla_finetuned \
  --job_name=smolvla_multi_gpu \
  --policy.device=cuda \
  --training.learning_rate=1e-4 \
  --wandb.enable=true

内存优化配置

# 针对显存较小的 GPU
lerobot-train \
  --policy.path=lerobot/smolvla_base \
  --dataset.repo_id=${HF_USER}/my_dataset \
  --batch_size=16 \
  --steps=30000 \
  --output_dir=outputs/train/smolvla_finetuned \
  --job_name=smolvla_memory_optimized \
  --policy.device=cuda \
  --training.learning_rate=5e-5 \
  --training.gradient_accumulation_steps=4 \
  --training.mixed_precision=fp16 \
  --wandb.enable=true

参数详细说明

核心参数

参数	含义	推荐值	说明
--policy.path	预训练模型路径	lerobot/smolvla_base	HuggingFace 上的官方预训练模型
--dataset.repo_id	数据集仓库 ID	${HF_USER}/my_dataset	你的 HuggingFace 数据集
--dataset.root	数据集存储位置	/data/my_dataset	指定从本地目录读取，不从huggingface下载数据集（可选）
--batch_size	批处理大小	64	根据显存调整，RTX 3080 推荐 32-64
--steps	训练步数	20000	小数据集可减少到 10000
--output_dir	输出目录	outputs/train/smolvla_finetuned	模型保存路径
--job_name	任务名称	smolvla_finetuning	用于日志和实验跟踪

训练参数

参数	含义	推荐值	说明
--training.learning_rate	学习率	1e-4	微调时可适当降低
--training.warmup_steps	预热步数	1000	学习率预热，稳定训练
--training.gradient_accumulation_steps	梯度累积步数	1	显存不足时可增加
--training.mixed_precision	混合精度	fp16	节省显存，加速训练
--training.max_grad_norm	梯度裁剪	1.0	防止梯度爆炸

模型特定参数

参数	含义	推荐值	说明
--policy.vision_backbone	视觉骨干网络	siglip	SmolVLA 默认使用 SigLIP
--policy.llm_backbone	语言模型骨干	qwen2	SmolVLA 默认使用 Qwen2
--policy.action_dim	动作维度	7	根据机器人自由度调整
--policy.chunk_size	动作块大小	50	预测的动作序列长度

训练监控

Weights & Biases 集成

SmolVLA 支持 W&B 进行实验跟踪：

# 启用 W&B 日志记录
lerobot-train \
  --wandb.enable=true \
  --wandb.project=smolvla_experiments \
  --wandb.run_name=smolvla_v1 \
  --wandb.notes="SmolVLA finetuning on custom dataset" \
  # ... 其他参数

关键指标监控

训练过程中需要关注的指标：

• Loss：总体损失，应该稳定下降
• Action Loss：动作预测损失
• Vision Loss：视觉特征损失
• Language Loss：语言理解损失
• Learning Rate：学习率变化
• GPU Memory：显存使用情况

模型评估

保存和加载模型

# 加载微调后的模型
from lerobot.common.policies.smolvla.modeling_smolvla import SmolVLAPolicy

policy = SmolVLAPolicy.from_pretrained(
    "outputs/train/smolvla_finetuned/checkpoints/last",
    device="cuda"
)

# 进行推理
observation = {
    "observation.images.cam_high": image_tensor,
    "observation.state": state_tensor
}

action = policy.select_action(observation)

性能评估脚本

# evaluation.py
import torch
from lerobot.common.policies.smolvla.modeling_smolvla import SmolVLAPolicy
from lerobot.common.datasets.lerobot_dataset import LeRobotDataset

def evaluate_model(model_path, dataset_path):
    # 加载模型
    policy = SmolVLAPolicy.from_pretrained(model_path, device="cuda")
    
    # 加载测试数据集
    dataset = LeRobotDataset(dataset_path, split="test")
    
    total_loss = 0
    num_samples = 0
    
    with torch.no_grad():
        for batch in dataset:
            prediction = policy(batch)
            loss = policy.compute_loss(prediction, batch)
            total_loss += loss.item()
            num_samples += 1
    
    avg_loss = total_loss / num_samples
    print(f"Average test loss: {avg_loss:.4f}")
    
    return avg_loss

if __name__ == "__main__":
    model_path = "outputs/train/smolvla_finetuned/checkpoints/last"
    dataset_path = "path/to/your/test/dataset"
    evaluate_model(model_path, dataset_path)

部署和推理

实时推理示例

# inference.py
import torch
import numpy as np
from PIL import Image
from lerobot.common.policies.smolvla.modeling_smolvla import SmolVLAPolicy

class SmolVLAInference:
    def __init__(self, model_path):
        self.policy = SmolVLAPolicy.from_pretrained(
            model_path, 
            device="cuda"
        )
        self.policy.eval()
    
    def predict_action(self, image, state, instruction=""):
        # 预处理图像
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image)
        
        # 构建观察
        observation = {
            "observation.images.cam_high": self.preprocess_image(image),
            "observation.state": torch.tensor(state, dtype=torch.float32).unsqueeze(0),
            "task.language_instruction": instruction
        }
        
        # 预测动作
        with torch.no_grad():
            action = self.policy.select_action(observation)
        
        return action.cpu().numpy()
    
    def preprocess_image(self, image):
        # 图像预处理逻辑
        image = image.resize((224, 224))
        image_tensor = torch.tensor(np.array(image)).permute(2, 0, 1).float() / 255.0
        return image_tensor.unsqueeze(0)

# 使用示例
if __name__ == "__main__":
    inference = SmolVLAInference("outputs/train/smolvla_finetuned/checkpoints/last")
    
    # 模拟输入
    image = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
    state = np.random.randn(7)  # 7-DOF 机器人状态
    instruction = "pick up the red cube"
    
    action = inference.predict_action(image, state, instruction)
    print(f"Predicted action: {action}")

最佳实践

数据准备建议

1. 数据质量：确保演示数据的质量，避免错误或不一致的动作
2. 数据多样性：包含不同场景、光照条件和物体位置的数据
3. 任务描述：为每个 episode 提供清晰的自然语言描述
4. 数据平衡：确保成功和失败案例的平衡

训练优化建议

1. 学习率调度：使用学习率预热和衰减策略
2. 正则化：适当使用 dropout 和权重衰减
3. 检查点保存：定期保存模型检查点
4. 早停策略：监控验证损失，避免过拟合

硬件优化建议

1. 显存管理：使用混合精度训练节省显存
2. 批处理大小：根据显存大小调整批处理大小
3. 数据加载：使用多进程数据加载加速训练
4. 模型并行：对于大模型，考虑使用模型并行

常见问题 (FAQ)

Q: SmolVLA 与其他 VLA 模型相比有什么优势？

A: SmolVLA 的主要优势包括：

• 轻量级：仅 450M 参数，适合消费级硬件
• 高效训练：训练时间相对较短
• 良好性能：在多个机器人任务上表现优异
• 易于部署：模型大小适中，便于实际部署

Q: 训练需要多长时间？

A: 训练时间取决于多个因素：

• 数据集大小：100 episodes 约需 2-4 小时（RTX 3080）
• 批处理大小：较大批处理可加速训练
• 硬件配置：更好的 GPU 可显著减少训练时间
• 训练步数：20000 步通常足够获得良好效果

Q: 如何判断模型是否收敛？

A: 观察以下指标：

• 损失曲线：总损失应稳定下降并趋于平稳
• 验证性能：验证集上的性能不再提升
• 动作预测：模型预测的动作应该合理
• 实际测试：在真实环境中测试模型性能

Q: 显存不足怎么办？

A: 可以尝试以下方法：

• 减少批处理大小（如从 64 减到 32 或 16）
• 启用混合精度训练（--training.mixed_precision=fp16）
• 增加梯度累积步数
• 使用更小的图像分辨率
• 考虑使用 CPU 卸载技术

Q: 如何提升模型性能？

A: 提升性能的方法：

• 增加数据量：收集更多高质量的演示数据
• 数据增强：使用图像增强技术增加数据多样性
• 超参数调优：调整学习率、批处理大小等参数
• 模型集成：训练多个模型并进行集成
• 领域适应：针对特定任务进行额外的微调

相关资源

• SmolVLA 官方博客
• LeRobot 官方文档
• SmolVLA 模型页面
• LeRobot GitHub 仓库
• 机器人学习论文集

更新日志

• 2024-01: 初始版本发布
• 2024-02: 添加多 GPU 训练支持
• 2024-03: 优化内存使用和训练效率
• 2024-04: 添加更多评估和部署示例