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.type smolvla \
  --policy.pretrained_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 \
  --policy.optimizer_lr 1e-4 \
  --policy.scheduler_warmup_steps 1000 \
  --policy.push_to_hub false \
  --save_checkpoint true \
  --save_freq 5000 \
  --wandb.enable true \
  --wandb.project smolvla_finetuning

高级训练配置

多 GPU 训练

# 使用 torchrun 进行多 GPU 训练
torchrun --nproc_per_node=2 --master_port=29500 \
  $(which lerobot-train) \
  --policy.type smolvla \
  --policy.pretrained_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 \
  --policy.optimizer_lr 1e-4 \
  --policy.push_to_hub false \
  --save_checkpoint true \
  --wandb.enable true

内存优化配置

# 针对显存较小的 GPU
lerobot-train \
  --policy.type smolvla \
  --policy.pretrained_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 \
  --policy.optimizer_lr 5e-5 \
  --policy.use_amp true \
  --num_workers 2 \
  --policy.push_to_hub false \
  --save_checkpoint true \
  --wandb.enable true

参数详细说明

核心参数

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

训练参数

参数	含义	推荐值	说明
`--policy.optimizer_lr`	学习率	`1e-4`	微调时可适当降低
`--policy.scheduler_warmup_steps`	预热步数	`1000`	学习率预热，稳定训练
`--policy.use_amp`	混合精度	`true`	节省显存，加速训练
`--policy.optimizer_grad_clip_norm`	梯度裁剪	`1.0`	防止梯度爆炸
`--num_workers`	数据加载线程数	`4`	根据CPU核心数调整
`--policy.push_to_hub`	推送到Hub	`false`	是否上传模型到HuggingFace（需要repo_id）
`--save_checkpoint`	保存检查点	`true`	是否保存训练检查点
`--save_freq`	保存频率	`5000`	检查点保存间隔步数

模型特定参数

参数	含义	推荐值	说明
`--policy.vlm_model_name`	VLM骨干模型	`HuggingFaceTB/SmolVLM2-500M-Video-Instruct`	SmolVLA 使用的视觉语言模型
`--policy.chunk_size`	动作块大小	`50`	预测的动作序列长度
`--policy.n_action_steps`	执行动作步数	`50`	每次实际执行的动作数
`--policy.n_obs_steps`	观测历史步数	`1`	使用的历史观测帧数

训练监控

Weights & Biases 集成

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

# 启用 W&B 日志记录
lerobot-train \
  --policy.type smolvla \
  --dataset.repo_id your-name/your-repo \
  --batch_size 64 \
  --steps 20000 \
  --policy.push_to_hub false \
  --wandb.enable true \
  --wandb.project smolvla_experiments \
  --wandb.notes "SmolVLA finetuning on custom dataset" \
  # ... 其他参数

关键指标监控

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

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

模型评估

保存和加载模型

# 加载微调后的模型
from lerobot.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.policies.smolvla.modeling_smolvla import SmolVLAPolicy
from lerobot.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.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}")

最佳实践

数据准备建议

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

训练优化建议

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

硬件优化建议

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

常见问题 (FAQ)

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

A: SmolVLA 的主要优势包括：

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

Q: 训练需要多长时间？

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

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

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

A: 观察以下指标：

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

Q: 显存不足怎么办？

A: 可以尝试以下方法：

减少批处理大小（如从 64 减到 32 或 16）：--batch_size 16
启用混合精度训练：--policy.use_amp true
减少数据加载线程：--num_workers 2
使用更小的图像分辨率：--policy.resize_imgs_with_padding 224 224
减少观测步数：--policy.n_obs_steps 1

Q: 如何提升模型性能？

A: 提升性能的方法：

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

更新日志

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

概述​

先决条件​

系统要求​

环境准备​

1. 安装 LeRobot​

2. 安装额外依赖​

数据准备​

LeRobot 格式数据​

数据质量要求​

微调训练​

基本训练命令​

高级训练配置​

多 GPU 训练​

内存优化配置​

参数详细说明​

核心参数​

训练参数​

模型特定参数​

训练监控​

Weights & Biases 集成​

关键指标监控​

模型评估​

保存和加载模型​

性能评估脚本​

部署和推理​

实时推理示例​

最佳实践​

数据准备建议​

训练优化建议​

硬件优化建议​

常见问题 (FAQ)​

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

Q: 训练需要多长时间？​

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

Q: 显存不足怎么办？​

Q: 如何提升模型性能？​

相关资源​

更新日志​

概述