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:
qing committed
1 parent
37206daf13
commit
45e019952e
6 files changed
+394
-12
No files matched your search
@@ -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()
|
||||
Reference in new issue
Block a user