解决无法在Lerobot-Seed分支上多GPU训练问题
由于我购买的机械臂不是So-Arm101系列,而是由SeedStudio公司开发的STarAI机械臂,这款机械臂也是基于lerobot开发的。目前Lerobot主仓库是可以支持多GPU训练的,在官网https://huggingface.co/docs/lerobot/index中可以查看:

我的训练指令如下:
accelerate launch --num_processes=2 $(which lerobot-train) \
--dataset.repo_id=/mnt/sdb/dataset/zzh/ACT/record-test-new1 \
--policy.type=act \
--policy.device=cuda \
--policy.chunk_size=100 \
--policy.n_action_steps=50 \
--policy.use_amp=true \
--policy.repo_id=starai/my_policy \
--batch_size=32 \
--policy.optimizer_lr=2e-05 \
--num_workers=4 \
--output_dir=outputs/train/act_viola_test1 \
--job_name=act_viola_test \
--wandb.enable=False \
--steps=100000 \
--save_checkpoint=True \
--save_freq=50000
但是存在报错,排查了一下原因发现报错和多 GPU 本身无关,而是配置和 ACTPolicy 构造函数不匹配(不太清除跟其他policy有没有产生冲突,我还没做到那一步)
原因是在lerobot-main/src/lerobot/policies/factory.py中的476行传递了2个参数给策略类

而我们翻看lerobot-main/src/lerobot/policies/act/modeling_act.py中的__init__函数时:

发现ACTPolicy.init()确实只接受一个 config参数,不接受这两个额外参数。
我推测原因是因为Seed分支落后与Lerobot主分支,相关的代码还没来得及适配,我询问了技术支持,证实了这个想法
解决办法:
很简单,既然factory传入了两个额外的参数,那我只需要在ACTPolicy.init()函数后加入2个形参,然后不去管它就好了
将函数头修改为:
def __init__(
self,
config: ACTConfig,
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
dataset_meta: Any | None = None,
):
并且导入库
from typing import Any
再次启动训练指令,嗯!错误消失,可以愉快的训练了。


浙公网安备 33010602011771号