actpred: action_dim改18 + .mt脚本预训练入口
- config.py: action_dim默认值 8→18(对齐bot4cj 18列) - mt_parser.py: 轻量级.mt @Action解析器(支持-占位符前向填充) - pretrain_mt.py: 预训练入口(零图像占位,checkpoint兼容现有推理) - _run.sh: 新增pretrain子命令 - README.md: 新增预训练文档 + action_dim说明 - test_actpred.py: 测试断言更新为18
This commit is contained in:
1 parent
37206daf13
commit
45e019952e
6 files changed
+394
-12
No files matched your search
+24
-4
@@ -21,11 +21,11 @@
|
|||||||
cd actpred && bash _set.sh
|
cd actpred && bash _set.sh
|
||||||
|
|
||||||
# 训练
|
# 训练
|
||||||
bash _run.sh train --data-dir ./data --epochs 100 --action-dim 8
|
bash _run.sh train --data-dir ./data --epochs 100 --action-dim 18
|
||||||
|
|
||||||
# 推理(--action-history 格式:冒号分隔各步,每步 action_dim 个值,共 history_len 步)
|
# 推理(--action-history 格式:冒号分隔各步,每步 action_dim 个值,共 history_len 步)
|
||||||
bash _run.sh --model checkpoints/best_inference.pt --image frame.jpg \
|
bash _run.sh --model checkpoints/best_inference.pt --image frame.jpg \
|
||||||
--action-history "0,0,0,0,0,0,0,0:1,0,0,0,0,0,0,0:0,1,0,0,0,0,0,0:0,0,1,0,0,0,0,0:0,0,0,1,0,0,0,0"
|
--action-history "45,0,45,0,45,0,45,0,90,90,10,200,0,0,0,0,0,0:45,0,45,0,45,0,45,0,90,90,0,20,0,0,0,0,0,0:45,0,45,0,45,0,45,0,90,90,0,0,0,0,0,0,0,0:45,0,45,0,45,0,45,0,90,90,0,0,0,0,0,0,0,0:45,0,45,0,45,0,45,0,90,90,0,0,0,0,0,0,0,0"
|
||||||
|
|
||||||
# 导出 ONNX(如安装了 onnxruntime,导出后自动验证一致性)
|
# 导出 ONNX(如安装了 onnxruntime,导出后自动验证一致性)
|
||||||
bash _run.sh export --model checkpoints/best_inference.pt --output model.onnx
|
bash _run.sh export --model checkpoints/best_inference.pt --output model.onnx
|
||||||
@@ -34,6 +34,26 @@ bash _run.sh export --model checkpoints/best_inference.pt --output model.onnx
|
|||||||
python test_actpred.py
|
python test_actpred.py
|
||||||
```
|
```
|
||||||
|
|
||||||
|
## 预训练(从 .mt 脚本,无需图像)
|
||||||
|
|
||||||
|
从 bot4cj 的 `.mt` 动作脚本中提取 `@Action` 矩阵,用零图像占位预训练动作动力学(history → next action),
|
||||||
|
保存的 checkpoint 与正常训练格式完全兼容,可直接用于推理或作为微调起点。
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# 预训练(解析 ../bot4cj/scripts/*.mt)
|
||||||
|
bash _run.sh pretrain --mt-dir ../bot4cj/scripts --epochs 200 --noise-std 1.0
|
||||||
|
|
||||||
|
# 用预训练权重作为微调起点(需修改 train.py 加载逻辑,或手动加载)
|
||||||
|
bash _run.sh train --data-dir ./data --epochs 50
|
||||||
|
```
|
||||||
|
|
||||||
|
| 参数 | 说明 |
|
||||||
|
|------|------|
|
||||||
|
| `--mt-dir` | .mt 脚本目录(如 `../bot4cj/scripts`) |
|
||||||
|
| `--noise-std` | 高斯噪声标准差(数据增强,建议 0.5~2.0) |
|
||||||
|
|
||||||
|
> **注意**:.mt 脚本列数(12)与 `action_dim`(18)不一致时自动补零。预训练保存到 `checkpoints/pretrain_best_inference.pt`。
|
||||||
|
|
||||||
## 数据格式
|
## 数据格式
|
||||||
|
|
||||||
```
|
```
|
||||||
@@ -55,7 +75,7 @@ data/
|
|||||||
```python
|
```python
|
||||||
from dataset import DataRecorder
|
from dataset import DataRecorder
|
||||||
|
|
||||||
recorder = DataRecorder("./data", action_dim=8)
|
recorder = DataRecorder("./data", action_dim=18)
|
||||||
for image, action in robot_loop():
|
for image, action in robot_loop():
|
||||||
recorder.record(image, action)
|
recorder.record(image, action)
|
||||||
recorder.flush()
|
recorder.flush()
|
||||||
@@ -68,7 +88,7 @@ recorder.flush()
|
|||||||
| 参数 | 默认值 | 说明 |
|
| 参数 | 默认值 | 说明 |
|
||||||
|------|--------|------|
|
|------|--------|------|
|
||||||
| `--image-size` | 64 | 图像缩放尺寸 |
|
| `--image-size` | 64 | 图像缩放尺寸 |
|
||||||
| `--action-dim` | 8 | 动作向量维度 |
|
| `--action-dim` | 18 | 动作向量维度(对齐 bot4cj 18 列:10 舵机 + 2 云台 + 扬声器频率 + 延时 + 4 预留)|
|
||||||
| `--history-len` | 5 | 历史动作窗口长度 |
|
| `--history-len` | 5 | 历史动作窗口长度 |
|
||||||
| `--backbone` | simple | CNN 类型 (simple/mobilenet) |
|
| `--backbone` | simple | CNN 类型 (simple/mobilenet) |
|
||||||
| `--history-encoder` | mlp | 历史编码器 (mlp/gru) |
|
| `--history-encoder` | mlp | 历史编码器 (mlp/gru) |
|
||||||
|
|||||||
@@ -5,6 +5,9 @@ source ./venv/bin/activate
|
|||||||
if [ "${1:-}" = "train" ]; then
|
if [ "${1:-}" = "train" ]; then
|
||||||
shift
|
shift
|
||||||
exec python3 train.py "$@"
|
exec python3 train.py "$@"
|
||||||
|
elif [ "${1:-}" = "pretrain" ]; then
|
||||||
|
shift
|
||||||
|
exec python3 pretrain_mt.py "$@"
|
||||||
elif [ "${1:-}" = "export" ]; then
|
elif [ "${1:-}" = "export" ]; then
|
||||||
shift
|
shift
|
||||||
exec python3 export_onnx.py "$@"
|
exec python3 export_onnx.py "$@"
|
||||||
|
|||||||
+1
-1
@@ -6,7 +6,7 @@ from typing import Optional
|
|||||||
@dataclass
|
@dataclass
|
||||||
class DataConfig:
|
class DataConfig:
|
||||||
image_size: int = 64
|
image_size: int = 64
|
||||||
action_dim: int = 8
|
action_dim: int = 18
|
||||||
history_len: int = 5
|
history_len: int = 5
|
||||||
data_dir: str = "./data"
|
data_dir: str = "./data"
|
||||||
normalize_actions: bool = False
|
normalize_actions: bool = False
|
||||||
|
|||||||
@@ -0,0 +1,91 @@
|
|||||||
|
"""轻量级 .mt 脚本 @Action 矩阵解析器
|
||||||
|
|
||||||
|
从 .mt 文件中提取 @Action 块,解析为 (n_rows, n_cols) 的 float32 矩阵。
|
||||||
|
支持 `-` 占位符(沿用上一行同列的值)。
|
||||||
|
"""
|
||||||
|
|
||||||
|
import re
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
|
||||||
|
def _natural_key(p: Path):
|
||||||
|
"""自然排序键,确保 s1.mt < s2.mt < s10.mt"""
|
||||||
|
return [int(t) if t.isdigit() else t for t in re.split(r"(\d+)", p.name)]
|
||||||
|
|
||||||
|
|
||||||
|
def parse_mt_action(path) -> np.ndarray:
|
||||||
|
"""解析 .mt 文件中的 @Action 矩阵
|
||||||
|
|
||||||
|
Args:
|
||||||
|
path: .mt 文件路径
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
(n_data_rows, n_cols) float32 矩阵,表头行已跳过,`-` 已前向填充
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: 文件中无 @Action 块或数据行不足
|
||||||
|
"""
|
||||||
|
text = Path(path).read_text(encoding="utf-8")
|
||||||
|
# 先去除整行注释(# 开头),避免注释中的 @Action 干扰正则匹配
|
||||||
|
text = re.sub(r"(?m)^[ \t]*#.*$", "", text)
|
||||||
|
|
||||||
|
match = re.search(r"@Action\s+(.*?)\s*@", text, re.DOTALL)
|
||||||
|
if not match:
|
||||||
|
raise ValueError(f"未找到 @Action 块: {path}")
|
||||||
|
|
||||||
|
block = match.group(1)
|
||||||
|
|
||||||
|
rows = []
|
||||||
|
for line in block.split("\n"):
|
||||||
|
line = re.sub(r"#.*$", "", line).strip()
|
||||||
|
if not line:
|
||||||
|
continue
|
||||||
|
if line.endswith(";"):
|
||||||
|
line = line[:-1].strip()
|
||||||
|
if not line:
|
||||||
|
continue
|
||||||
|
cols = re.split(r"\s+", line)
|
||||||
|
rows.append(cols)
|
||||||
|
|
||||||
|
if len(rows) < 2:
|
||||||
|
raise ValueError(
|
||||||
|
f"@Action 至少需要表头 + 1 行数据,实际 {len(rows)} 行: {path}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# 第一行是表头(列名如 s1.A, s2.A, ...),跳过
|
||||||
|
data_rows = rows[1:]
|
||||||
|
n_cols = len(rows[0])
|
||||||
|
|
||||||
|
matrix = np.zeros((len(data_rows), n_cols), dtype=np.float32)
|
||||||
|
for i, row in enumerate(data_rows):
|
||||||
|
for j, val in enumerate(row):
|
||||||
|
if val == "-":
|
||||||
|
matrix[i, j] = matrix[i - 1, j] if i > 0 else 0.0
|
||||||
|
else:
|
||||||
|
matrix[i, j] = float(val)
|
||||||
|
|
||||||
|
return matrix
|
||||||
|
|
||||||
|
|
||||||
|
def parse_mt_directory(mt_dir, pattern="*.mt"):
|
||||||
|
"""解析目录下所有 .mt 文件的 @Action 矩阵
|
||||||
|
|
||||||
|
Args:
|
||||||
|
mt_dir: .mt 文件目录
|
||||||
|
pattern: 文件匹配模式
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
list of (filepath, np.ndarray) 元组,按文件名自然排序
|
||||||
|
"""
|
||||||
|
dir_path = Path(mt_dir)
|
||||||
|
files = sorted(dir_path.glob(pattern), key=_natural_key)
|
||||||
|
results = []
|
||||||
|
for f in files:
|
||||||
|
try:
|
||||||
|
matrix = parse_mt_action(f)
|
||||||
|
results.append((str(f), matrix))
|
||||||
|
except ValueError:
|
||||||
|
pass
|
||||||
|
return results
|
||||||
@@ -0,0 +1,268 @@
|
|||||||
|
"""actpred 预训练入口:从 .mt 脚本学习动作动力学(无需图像)
|
||||||
|
|
||||||
|
用法:
|
||||||
|
bash _run.sh pretrain --mt-dir ../bot4cj/scripts --epochs 200
|
||||||
|
|
||||||
|
流程:
|
||||||
|
1. 解析 .mt 文件中的 @Action 矩阵 → 动作序列
|
||||||
|
2. 滑动窗口采样: action_history[H] → target action(下一行)
|
||||||
|
3. 用零图像占位训练完整 ActionPredictor(保持 checkpoint 兼容性)
|
||||||
|
4. 保存 checkpoint,格式与 train.py 一致,可直接用于推理或微调
|
||||||
|
|
||||||
|
数据增强:
|
||||||
|
--noise-std 控制高斯噪声标准差。因 .mt 脚本数据量有限,
|
||||||
|
适度噪声(如 0.5~2.0)可有效防止过拟合。
|
||||||
|
"""
|
||||||
|
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
import random
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import torch
|
||||||
|
from torch.utils.data import Dataset, DataLoader, Subset
|
||||||
|
from tqdm import tqdm
|
||||||
|
|
||||||
|
from config import Config, build_argparser, apply_overrides
|
||||||
|
from model import ActionPredictor
|
||||||
|
from mt_parser import parse_mt_directory
|
||||||
|
|
||||||
|
|
||||||
|
class MTActionDataset(Dataset):
|
||||||
|
"""从 .mt @Action 矩阵构建的滑动窗口数据集(无图像)
|
||||||
|
|
||||||
|
每个样本: (action_history[history_len, action_dim], target[action_dim])
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, mt_dir: str, history_len: int, action_dim: int):
|
||||||
|
self.history_len = history_len
|
||||||
|
self.action_dim = action_dim
|
||||||
|
self.samples = []
|
||||||
|
|
||||||
|
for path, matrix in parse_mt_directory(mt_dir):
|
||||||
|
n_rows = matrix.shape[1]
|
||||||
|
if n_rows < action_dim:
|
||||||
|
matrix = np.pad(
|
||||||
|
matrix,
|
||||||
|
((0, 0), (0, action_dim - n_rows)),
|
||||||
|
mode="constant",
|
||||||
|
)
|
||||||
|
elif n_rows > action_dim:
|
||||||
|
matrix = matrix[:, :action_dim]
|
||||||
|
|
||||||
|
for i in range(len(matrix) - history_len):
|
||||||
|
history = matrix[i : i + history_len]
|
||||||
|
target = matrix[i + history_len]
|
||||||
|
self.samples.append((history.copy(), target.copy()))
|
||||||
|
|
||||||
|
print(f"MTActionDataset: {len(self.samples)} samples "
|
||||||
|
f"from .mt files (history_len={history_len}, action_dim={action_dim})")
|
||||||
|
|
||||||
|
def __len__(self):
|
||||||
|
return len(self.samples)
|
||||||
|
|
||||||
|
def __getitem__(self, idx):
|
||||||
|
history, target = self.samples[idx]
|
||||||
|
return {
|
||||||
|
"action_history": torch.from_numpy(history.copy()),
|
||||||
|
"target": torch.from_numpy(target.copy()),
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def pretrain_one_epoch(model, loader, optimizer, device, image_size,
|
||||||
|
max_grad_norm=None, noise_std=0.0):
|
||||||
|
model.train()
|
||||||
|
total_loss = 0.0
|
||||||
|
count = 0
|
||||||
|
zero_image = None
|
||||||
|
for batch in tqdm(loader, desc=" pretrain", leave=False):
|
||||||
|
history = batch["action_history"].to(device)
|
||||||
|
target = batch["target"].to(device)
|
||||||
|
if noise_std > 0:
|
||||||
|
history = history + torch.randn_like(history) * noise_std
|
||||||
|
target = target + torch.randn_like(target) * noise_std
|
||||||
|
if zero_image is None or zero_image.size(0) != history.size(0):
|
||||||
|
zero_image = torch.zeros(
|
||||||
|
history.size(0), 3, image_size, image_size, device=device
|
||||||
|
)
|
||||||
|
optimizer.zero_grad()
|
||||||
|
pred = model(zero_image, history)
|
||||||
|
loss = torch.nn.functional.mse_loss(pred, target)
|
||||||
|
loss.backward()
|
||||||
|
if max_grad_norm:
|
||||||
|
torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
|
||||||
|
optimizer.step()
|
||||||
|
total_loss += loss.item() * history.size(0)
|
||||||
|
count += history.size(0)
|
||||||
|
return total_loss / max(count, 1)
|
||||||
|
|
||||||
|
|
||||||
|
@torch.no_grad()
|
||||||
|
def validate(model, loader, device, image_size):
|
||||||
|
model.eval()
|
||||||
|
total_loss = 0.0
|
||||||
|
count = 0
|
||||||
|
zero_image = None
|
||||||
|
for batch in tqdm(loader, desc=" val ", leave=False):
|
||||||
|
history = batch["action_history"].to(device)
|
||||||
|
target = batch["target"].to(device)
|
||||||
|
if zero_image is None or zero_image.size(0) != history.size(0):
|
||||||
|
zero_image = torch.zeros(
|
||||||
|
history.size(0), 3, image_size, image_size, device=device
|
||||||
|
)
|
||||||
|
pred = model(zero_image, history)
|
||||||
|
loss = torch.nn.functional.mse_loss(pred, target)
|
||||||
|
total_loss += loss.item() * history.size(0)
|
||||||
|
count += history.size(0)
|
||||||
|
if count == 0:
|
||||||
|
return float("inf")
|
||||||
|
return total_loss / count
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
parser = build_argparser("actpred pre-training from .mt scripts")
|
||||||
|
parser.add_argument("--mt-dir", type=str, required=True,
|
||||||
|
help=".mt 脚本目录(如 ../bot4cj/scripts)")
|
||||||
|
parser.add_argument("--noise-std", type=float, default=0.0,
|
||||||
|
help="高斯噪声标准差(数据增强,建议 0.5~2.0)")
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
cfg = apply_overrides(Config(), args)
|
||||||
|
cfg.data.data_dir = args.mt_dir
|
||||||
|
|
||||||
|
os.makedirs(cfg.train.checkpoint_dir, exist_ok=True)
|
||||||
|
|
||||||
|
random.seed(cfg.train.seed)
|
||||||
|
np.random.seed(cfg.train.seed)
|
||||||
|
torch.manual_seed(cfg.train.seed)
|
||||||
|
if torch.cuda.is_available():
|
||||||
|
torch.cuda.manual_seed_all(cfg.train.seed)
|
||||||
|
|
||||||
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||||
|
print(f"Device: {device}")
|
||||||
|
|
||||||
|
full_ds = MTActionDataset(
|
||||||
|
args.mt_dir, cfg.data.history_len, cfg.data.action_dim
|
||||||
|
)
|
||||||
|
if len(full_ds) == 0:
|
||||||
|
raise ValueError(
|
||||||
|
f"数据集为空。请检查 mt-dir={args.mt_dir} "
|
||||||
|
f"及 .mt 文件中 @Action 行数 > history_len={cfg.data.history_len}"
|
||||||
|
)
|
||||||
|
|
||||||
|
n = len(full_ds)
|
||||||
|
indices = list(range(n))
|
||||||
|
rng = np.random.RandomState(cfg.train.seed)
|
||||||
|
rng.shuffle(indices)
|
||||||
|
val_count = max(1, int(n * cfg.train.val_split))
|
||||||
|
val_indices = indices[:val_count]
|
||||||
|
train_indices = indices[val_count:]
|
||||||
|
|
||||||
|
train_ds = Subset(full_ds, train_indices)
|
||||||
|
val_ds = Subset(full_ds, val_indices)
|
||||||
|
|
||||||
|
print(f"Samples: {len(train_ds)} train, {len(val_ds)} val")
|
||||||
|
if len(train_ds) == 0:
|
||||||
|
raise ValueError("训练集为空(样本数不足)")
|
||||||
|
if len(val_ds) == 0:
|
||||||
|
print("WARNING: 验证集为空,模型选择禁用")
|
||||||
|
|
||||||
|
loader_gen = torch.Generator()
|
||||||
|
loader_gen.manual_seed(cfg.train.seed)
|
||||||
|
train_loader = DataLoader(
|
||||||
|
train_ds, batch_size=cfg.train.batch_size,
|
||||||
|
shuffle=True, num_workers=cfg.train.num_workers,
|
||||||
|
generator=loader_gen,
|
||||||
|
)
|
||||||
|
val_loader = DataLoader(
|
||||||
|
val_ds, batch_size=cfg.train.batch_size,
|
||||||
|
shuffle=False, num_workers=cfg.train.num_workers,
|
||||||
|
)
|
||||||
|
|
||||||
|
model = ActionPredictor(cfg.model, cfg.data).to(device)
|
||||||
|
n_params = sum(p.numel() for p in model.parameters())
|
||||||
|
print(f"Model params: {n_params:,}")
|
||||||
|
print(f"Config: action_dim={cfg.data.action_dim}, "
|
||||||
|
f"history_len={cfg.data.history_len}, "
|
||||||
|
f"backbone={cfg.model.cnn_backbone}, "
|
||||||
|
f"encoder={cfg.model.history_encoder}")
|
||||||
|
print(f"Pre-training: noise_std={args.noise_std}")
|
||||||
|
|
||||||
|
optimizer = torch.optim.AdamW(
|
||||||
|
model.parameters(), lr=cfg.train.lr,
|
||||||
|
weight_decay=cfg.train.weight_decay,
|
||||||
|
)
|
||||||
|
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
|
||||||
|
optimizer, T_max=cfg.train.epochs
|
||||||
|
)
|
||||||
|
|
||||||
|
best_val_loss = float("inf")
|
||||||
|
patience_counter = 0
|
||||||
|
best_path = os.path.join(cfg.train.checkpoint_dir, "pretrain_best.pt")
|
||||||
|
inference_path = os.path.join(
|
||||||
|
cfg.train.checkpoint_dir, "pretrain_best_inference.pt"
|
||||||
|
)
|
||||||
|
|
||||||
|
for epoch in range(1, cfg.train.epochs + 1):
|
||||||
|
print(f"\nEpoch {epoch}/{cfg.train.epochs} "
|
||||||
|
f"lr={optimizer.param_groups[0]['lr']:.6f}")
|
||||||
|
train_loss = pretrain_one_epoch(
|
||||||
|
model, train_loader, optimizer, device,
|
||||||
|
cfg.data.image_size, cfg.train.max_grad_norm, args.noise_std,
|
||||||
|
)
|
||||||
|
val_loss = validate(model, val_loader, device, cfg.data.image_size)
|
||||||
|
scheduler.step()
|
||||||
|
print(f" train_loss={train_loss:.6f} val_loss={val_loss:.6f}")
|
||||||
|
|
||||||
|
if val_loss < best_val_loss:
|
||||||
|
best_val_loss = val_loss
|
||||||
|
patience_counter = 0
|
||||||
|
torch.save(
|
||||||
|
{
|
||||||
|
"epoch": epoch,
|
||||||
|
"model_state_dict": model.state_dict(),
|
||||||
|
"optimizer_state_dict": optimizer.state_dict(),
|
||||||
|
"val_loss": val_loss,
|
||||||
|
"config": cfg.to_dict(),
|
||||||
|
},
|
||||||
|
best_path,
|
||||||
|
)
|
||||||
|
torch.save(
|
||||||
|
{
|
||||||
|
"model_state_dict": model.state_dict(),
|
||||||
|
"config": cfg.to_dict(),
|
||||||
|
},
|
||||||
|
inference_path,
|
||||||
|
)
|
||||||
|
print(f" -> saved best (val_loss={val_loss:.6f})")
|
||||||
|
else:
|
||||||
|
patience_counter += 1
|
||||||
|
if patience_counter >= cfg.train.patience:
|
||||||
|
print(f"\nEarly stopping at epoch {epoch} "
|
||||||
|
f"(patience={cfg.train.patience})")
|
||||||
|
break
|
||||||
|
|
||||||
|
print(f"\nPre-training done. Best val_loss={best_val_loss:.6f}")
|
||||||
|
|
||||||
|
if best_val_loss == float("inf"):
|
||||||
|
torch.save(
|
||||||
|
{
|
||||||
|
"model_state_dict": model.state_dict(),
|
||||||
|
"config": cfg.to_dict(),
|
||||||
|
},
|
||||||
|
inference_path,
|
||||||
|
)
|
||||||
|
print("WARNING: saved final weights (no valid validation loss)")
|
||||||
|
else:
|
||||||
|
print(f"Best pre-trained model: {inference_path}")
|
||||||
|
print("Fine-tune: bash _run.sh train --data-dir ./data "
|
||||||
|
f"--model {inference_path}")
|
||||||
|
|
||||||
|
if args.save_config:
|
||||||
|
with open(args.save_config, "w") as f:
|
||||||
|
json.dump(cfg.to_dict(), f, indent=2, ensure_ascii=False)
|
||||||
|
print(f"Config saved to: {args.save_config}")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@@ -60,28 +60,28 @@ class TestActionPredictor(unittest.TestCase):
|
|||||||
"""测试模型前向传播"""
|
"""测试模型前向传播"""
|
||||||
|
|
||||||
def test_forward_simple_mlp(self):
|
def test_forward_simple_mlp(self):
|
||||||
data_cfg = DataConfig(image_size=64, action_dim=8, history_len=5)
|
data_cfg = DataConfig(image_size=64, action_dim=18, history_len=5)
|
||||||
model_cfg = ModelConfig(cnn_backbone="simple", history_encoder="mlp")
|
model_cfg = ModelConfig(cnn_backbone="simple", history_encoder="mlp")
|
||||||
model = ActionPredictor(model_cfg, data_cfg)
|
model = ActionPredictor(model_cfg, data_cfg)
|
||||||
model.eval()
|
model.eval()
|
||||||
|
|
||||||
image = torch.randn(2, 3, 64, 64)
|
image = torch.randn(2, 3, 64, 64)
|
||||||
history = torch.randn(2, 5, 8)
|
history = torch.randn(2, 5, 18)
|
||||||
with torch.no_grad():
|
with torch.no_grad():
|
||||||
out = model(image, history)
|
out = model(image, history)
|
||||||
self.assertEqual(out.shape, (2, 8))
|
self.assertEqual(out.shape, (2, 18))
|
||||||
|
|
||||||
def test_forward_simple_gru(self):
|
def test_forward_simple_gru(self):
|
||||||
data_cfg = DataConfig(image_size=64, action_dim=8, history_len=5)
|
data_cfg = DataConfig(image_size=64, action_dim=18, history_len=5)
|
||||||
model_cfg = ModelConfig(cnn_backbone="simple", history_encoder="gru")
|
model_cfg = ModelConfig(cnn_backbone="simple", history_encoder="gru")
|
||||||
model = ActionPredictor(model_cfg, data_cfg)
|
model = ActionPredictor(model_cfg, data_cfg)
|
||||||
model.eval()
|
model.eval()
|
||||||
|
|
||||||
image = torch.randn(2, 3, 64, 64)
|
image = torch.randn(2, 3, 64, 64)
|
||||||
history = torch.randn(2, 5, 8)
|
history = torch.randn(2, 5, 18)
|
||||||
with torch.no_grad():
|
with torch.no_grad():
|
||||||
out = model(image, history)
|
out = model(image, history)
|
||||||
self.assertEqual(out.shape, (2, 8))
|
self.assertEqual(out.shape, (2, 18))
|
||||||
|
|
||||||
def test_forward_dynamic_batch(self):
|
def test_forward_dynamic_batch(self):
|
||||||
data_cfg = DataConfig(image_size=32, action_dim=4, history_len=3)
|
data_cfg = DataConfig(image_size=32, action_dim=4, history_len=3)
|
||||||
@@ -186,7 +186,7 @@ class TestConfig(unittest.TestCase):
|
|||||||
def test_default_config(self):
|
def test_default_config(self):
|
||||||
cfg = Config()
|
cfg = Config()
|
||||||
self.assertEqual(cfg.data.image_size, 64)
|
self.assertEqual(cfg.data.image_size, 64)
|
||||||
self.assertEqual(cfg.data.action_dim, 8)
|
self.assertEqual(cfg.data.action_dim, 18)
|
||||||
self.assertEqual(cfg.data.history_len, 5)
|
self.assertEqual(cfg.data.history_len, 5)
|
||||||
self.assertEqual(cfg.model.cnn_backbone, "simple")
|
self.assertEqual(cfg.model.cnn_backbone, "simple")
|
||||||
self.assertEqual(cfg.model.history_encoder, "mlp")
|
self.assertEqual(cfg.model.history_encoder, "mlp")
|
||||||
|
|||||||
Reference in new issue
Block a user