Views
No views yet
100000/
├── pretrained_model/
│ ├── model.safetensors # 模型权重 (~1005MB)
│ ├── config.json # 模型配置
│ └── train_config.json # 训练配置
└── training_state/
├── optimizer_state.safetensors # 优化器状态 (~2GB)
├── scheduler_state.json # 学习率调度器状态
├── optimizer_param_groups.json # 优化器参数组
├── rng_state.safetensors # 随机数生成器状态
└── training_step.json # 训练步数信息1import torch
2from lerobot.common.policies.diffusion import DiffusionPolicy
3
4# 加载配置
5config_path = "pretrained_model/config.json"
6with open(config_path, 'r') as f:
7 config = json.load(f)
8
9# 创建模型
10policy = DiffusionPolicy(config)
11
12# 加载权重
13checkpoint = torch.load("pretrained_model/model.safetensors", map_location='cpu')
14policy.load_state_dict(checkpoint)1# 加载训练状态
2training_state = torch.load("training_state/optimizer_state.safetensors")
3optimizer.load_state_dict(training_state)
4
5# 继续训练
6# ... 训练代码