fix(actpred): 修复归一化统计不一致等 5 项问题并增强可复现性

- train/val 归一化统计量改为仅用训练集计算并落盘复用,修复 val loss 不可信
- predict.py 统计文件路径改用 os.path 处理,修复跨平台与无分隔符路径问题
- dataset.py jpg/png 混合时统一按文件名排序,保证时间顺序
- 新增 --seed 全局随机种子(含 DataLoader generator),训练可复现
- 新增 --max-grad-norm 梯度裁剪(默认 1.0),防 GRU 梯度爆炸
- get_episode_count 与数据集目录过滤条件对齐,避免无效目录导致索引错位
- 空验证集返回 inf 并保底保存最终权重,避免伪造 best 模型
- 清理 export_onnx.py 重复分支死代码
This commit is contained in:
qing committed 2026-08-08 14:00:49 +00:00
1 parent 85e5a4a054
commit 73b72d6889
5 files changed
+58 -18

No files matched your search

+6
View File
@@ -33,6 +33,8 @@ class TrainConfig:
checkpoint_dir: str = "./checkpoints" checkpoint_dir: str = "./checkpoints"
patience: int = 15 patience: int = 15
num_workers: int = 2 num_workers: int = 2
seed: int = 42
max_grad_norm: float = 1.0
@dataclass @dataclass
@@ -75,6 +77,8 @@ def build_argparser(description: str = "") -> argparse.ArgumentParser:
parser.add_argument("--checkpoint-dir", type=str, default=None) parser.add_argument("--checkpoint-dir", type=str, default=None)
parser.add_argument("--patience", type=int, default=None) parser.add_argument("--patience", type=int, default=None)
parser.add_argument("--num-workers", type=int, default=None) parser.add_argument("--num-workers", type=int, default=None)
parser.add_argument("--seed", type=int, default=None)
parser.add_argument("--max-grad-norm", type=float, default=None)
parser.add_argument("--save-config", type=str, default=None, help="保存配置到 JSON 文件") parser.add_argument("--save-config", type=str, default=None, help="保存配置到 JSON 文件")
return parser return parser
@@ -101,6 +105,8 @@ def apply_overrides(cfg: Config, args: argparse.Namespace) -> Config:
"checkpoint_dir": ("train", "checkpoint_dir"), "checkpoint_dir": ("train", "checkpoint_dir"),
"patience": ("train", "patience"), "patience": ("train", "patience"),
"num_workers": ("train", "num_workers"), "num_workers": ("train", "num_workers"),
"seed": ("train", "seed"),
"max_grad_norm": ("train", "max_grad_norm"),
} }
args_dict = vars(args) args_dict = vars(args)
for arg_name, (section, field) in mapping.items(): for arg_name, (section, field) in mapping.items():
+10 -3
View File
@@ -35,7 +35,8 @@ class ActionDataset(Dataset):
continue continue
actions = np.load(actions_path).astype(np.float32) actions = np.load(actions_path).astype(np.float32)
n_frames = len(actions) n_frames = len(actions)
image_files = sorted(images_dir.glob("*.jpg")) + sorted(images_dir.glob("*.png")) image_files = sorted(list(images_dir.glob("*.jpg")) + list(images_dir.glob("*.png")),
key=lambda p: p.name)
n_images = len(image_files) n_images = len(image_files)
n = min(n_frames, n_images) n = min(n_frames, n_images)
for t in range(history_len, n): for t in range(history_len, n):
@@ -73,9 +74,15 @@ class ActionDataset(Dataset):
return None return None
@staticmethod @staticmethod
def get_episode_count(data_dir: str) -> int: def _valid_episode_dirs(data_dir: str) -> list:
data_path = Path(data_dir) data_path = Path(data_dir)
return len([d for d in data_path.iterdir() if d.is_dir()]) return sorted([d for d in data_path.iterdir()
if d.is_dir() and (d / "actions.npy").exists() and (d / "images").is_dir()])
@staticmethod
def get_episode_count(data_dir: str) -> int:
# 与 __init__ 的过滤条件保持一致,避免 split_episodes 索引错位
return len(ActionDataset._valid_episode_dirs(data_dir))
class DataRecorder: class DataRecorder:
+1 -4
View File
@@ -20,10 +20,7 @@ def main():
model.eval() model.eval()
dummy_image = torch.randn(1, 3, data_cfg.image_size, data_cfg.image_size) dummy_image = torch.randn(1, 3, data_cfg.image_size, data_cfg.image_size)
if model_cfg.history_encoder == "gru": dummy_history = torch.randn(1, data_cfg.history_len, data_cfg.action_dim)
dummy_history = torch.randn(1, data_cfg.history_len, data_cfg.action_dim)
else:
dummy_history = torch.randn(1, data_cfg.history_len, data_cfg.action_dim)
torch.onnx.export( torch.onnx.export(
model, (dummy_image, dummy_history), model, (dummy_image, dummy_history),
+3 -1
View File
@@ -1,5 +1,6 @@
import argparse import argparse
import json import json
import os
import numpy as np import numpy as np
import torch import torch
from PIL import Image from PIL import Image
@@ -32,7 +33,8 @@ class ActionPredictorInference:
self._action_mean = None self._action_mean = None
self._action_std = None self._action_std = None
stats_path = checkpoint_path.rsplit("/", 1)[0] + "/action_stats.npy" stats_path = os.path.join(os.path.dirname(os.path.abspath(checkpoint_path)),
"action_stats.npy")
try: try:
stats = np.load(stats_path) stats = np.load(stats_path)
self._action_mean = stats[0] self._action_mean = stats[0]
+38 -10
View File
@@ -1,5 +1,6 @@
import json import json
import os import os
import random
import numpy as np import numpy as np
import torch import torch
from torch.utils.data import DataLoader, Subset from torch.utils.data import DataLoader, Subset
@@ -19,7 +20,7 @@ def split_episodes(data_dir: str, val_ratio: float, seed: int = 42):
return indices[val_count:], indices[:val_count] return indices[val_count:], indices[:val_count]
def train_one_epoch(model, loader, optimizer, device): def train_one_epoch(model, loader, optimizer, device, max_grad_norm: float = None):
model.train() model.train()
total_loss = 0.0 total_loss = 0.0
count = 0 count = 0
@@ -31,6 +32,8 @@ def train_one_epoch(model, loader, optimizer, device):
pred = model(image, history) pred = model(image, history)
loss = torch.nn.functional.mse_loss(pred, target) loss = torch.nn.functional.mse_loss(pred, target)
loss.backward() loss.backward()
if max_grad_norm:
torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
optimizer.step() optimizer.step()
total_loss += loss.item() * image.size(0) total_loss += loss.item() * image.size(0)
count += image.size(0) count += image.size(0)
@@ -50,7 +53,9 @@ def validate(model, loader, device):
loss = torch.nn.functional.mse_loss(pred, target) loss = torch.nn.functional.mse_loss(pred, target)
total_loss += loss.item() * image.size(0) total_loss += loss.item() * image.size(0)
count += image.size(0) count += image.size(0)
return total_loss / max(count, 1) if count == 0:
return float("inf") # 空验证集不参与比较,避免假 best
return total_loss / count
def main(): def main():
@@ -60,6 +65,13 @@ def main():
os.makedirs(cfg.train.checkpoint_dir, exist_ok=True) 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") device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device: {device}") print(f"Device: {device}")
@@ -71,15 +83,27 @@ def main():
cfg.data.data_dir, cfg.data.history_len, cfg.data.action_dim, cfg.data.data_dir, cfg.data.history_len, cfg.data.action_dim,
cfg.data.image_size, cfg.data.normalize_actions, cfg.data.image_size, cfg.data.normalize_actions,
cfg.data.action_stats_path, train_indices) cfg.data.action_stats_path, train_indices)
if cfg.data.normalize_actions and not cfg.data.action_stats_path:
# 归一化统计量必须只用训练集计算,再供验证集复用,否则 val loss 不可信
stats_path = os.path.join(cfg.train.checkpoint_dir, "action_stats.npy")
np.save(stats_path, train_ds.get_action_stats())
print(f"Action stats saved to: {stats_path}")
cfg.data.action_stats_path = stats_path
val_ds = ActionDataset( val_ds = ActionDataset(
cfg.data.data_dir, cfg.data.history_len, cfg.data.action_dim, cfg.data.data_dir, cfg.data.history_len, cfg.data.action_dim,
cfg.data.image_size, cfg.data.normalize_actions, cfg.data.image_size, cfg.data.normalize_actions,
cfg.data.action_stats_path, val_indices) cfg.data.action_stats_path, val_indices)
print(f"Samples: {len(train_ds)} train, {len(val_ds)} val") print(f"Samples: {len(train_ds)} train, {len(val_ds)} val")
if len(val_ds) == 0:
print("WARNING: validation set is empty (episodes too short/few); "
"model selection is disabled, last epoch will be saved.")
loader_gen = torch.Generator()
loader_gen.manual_seed(cfg.train.seed)
train_loader = DataLoader(train_ds, batch_size=cfg.train.batch_size, train_loader = DataLoader(train_ds, batch_size=cfg.train.batch_size,
shuffle=True, num_workers=cfg.train.num_workers) shuffle=True, num_workers=cfg.train.num_workers,
generator=loader_gen)
val_loader = DataLoader(val_ds, batch_size=cfg.train.batch_size, val_loader = DataLoader(val_ds, batch_size=cfg.train.batch_size,
shuffle=False, num_workers=cfg.train.num_workers) shuffle=False, num_workers=cfg.train.num_workers)
@@ -97,7 +121,8 @@ def main():
for epoch in range(1, cfg.train.epochs + 1): for epoch in range(1, cfg.train.epochs + 1):
print(f"\nEpoch {epoch}/{cfg.train.epochs} lr={optimizer.param_groups[0]['lr']:.6f}") print(f"\nEpoch {epoch}/{cfg.train.epochs} lr={optimizer.param_groups[0]['lr']:.6f}")
train_loss = train_one_epoch(model, train_loader, optimizer, device) train_loss = train_one_epoch(model, train_loader, optimizer, device,
cfg.train.max_grad_norm)
val_loss = validate(model, val_loader, device) val_loss = validate(model, val_loader, device)
scheduler.step() scheduler.step()
print(f" train_loss={train_loss:.6f} val_loss={val_loss:.6f}") print(f" train_loss={train_loss:.6f} val_loss={val_loss:.6f}")
@@ -125,13 +150,16 @@ def main():
break break
print(f"\nTraining done. Best val_loss={best_val_loss:.6f}") print(f"\nTraining done. Best val_loss={best_val_loss:.6f}")
print(f"Best model saved to: {best_path}")
action_stats = train_ds.get_action_stats() if best_val_loss == float("inf"):
if action_stats is not None: # 验证集为空时无 best 可选,保底保存最终权重
stats_path = os.path.join(cfg.train.checkpoint_dir, "action_stats.npy") torch.save({
np.save(stats_path, action_stats) "model_state_dict": model.state_dict(),
print(f"Action stats saved to: {stats_path}") "config": cfg.to_dict(),
}, os.path.join(cfg.train.checkpoint_dir, "best_inference.pt"))
print("WARNING: saved final weights as best_inference.pt (no valid validation loss)")
else:
print(f"Best model saved to: {best_path}")
if args.save_config: if args.save_config:
with open(args.save_config, "w") as f: with open(args.save_config, "w") as f: