diff --git a/actpred/config.py b/actpred/config.py index ee5f890..1bb3fad 100644 --- a/actpred/config.py +++ b/actpred/config.py @@ -33,6 +33,8 @@ class TrainConfig: checkpoint_dir: str = "./checkpoints" patience: int = 15 num_workers: int = 2 + seed: int = 42 + max_grad_norm: float = 1.0 @dataclass @@ -75,6 +77,8 @@ def build_argparser(description: str = "") -> argparse.ArgumentParser: parser.add_argument("--checkpoint-dir", type=str, default=None) parser.add_argument("--patience", 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 文件") return parser @@ -101,6 +105,8 @@ def apply_overrides(cfg: Config, args: argparse.Namespace) -> Config: "checkpoint_dir": ("train", "checkpoint_dir"), "patience": ("train", "patience"), "num_workers": ("train", "num_workers"), + "seed": ("train", "seed"), + "max_grad_norm": ("train", "max_grad_norm"), } args_dict = vars(args) for arg_name, (section, field) in mapping.items(): diff --git a/actpred/dataset.py b/actpred/dataset.py index 3a73d91..78a16cb 100644 --- a/actpred/dataset.py +++ b/actpred/dataset.py @@ -35,7 +35,8 @@ class ActionDataset(Dataset): continue actions = np.load(actions_path).astype(np.float32) 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 = min(n_frames, n_images) for t in range(history_len, n): @@ -73,9 +74,15 @@ class ActionDataset(Dataset): return None @staticmethod - def get_episode_count(data_dir: str) -> int: + def _valid_episode_dirs(data_dir: str) -> list: 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: diff --git a/actpred/export_onnx.py b/actpred/export_onnx.py index 1a19fb4..d3aebcf 100644 --- a/actpred/export_onnx.py +++ b/actpred/export_onnx.py @@ -20,10 +20,7 @@ def main(): model.eval() 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) - else: - 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) torch.onnx.export( model, (dummy_image, dummy_history), diff --git a/actpred/predict.py b/actpred/predict.py index 05d24ac..07a414e 100644 --- a/actpred/predict.py +++ b/actpred/predict.py @@ -1,5 +1,6 @@ import argparse import json +import os import numpy as np import torch from PIL import Image @@ -32,7 +33,8 @@ class ActionPredictorInference: self._action_mean = 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: stats = np.load(stats_path) self._action_mean = stats[0] diff --git a/actpred/train.py b/actpred/train.py index fc8cae4..2485675 100644 --- a/actpred/train.py +++ b/actpred/train.py @@ -1,5 +1,6 @@ import json import os +import random import numpy as np import torch 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] -def train_one_epoch(model, loader, optimizer, device): +def train_one_epoch(model, loader, optimizer, device, max_grad_norm: float = None): model.train() total_loss = 0.0 count = 0 @@ -31,6 +32,8 @@ def train_one_epoch(model, loader, optimizer, device): pred = model(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() * image.size(0) count += image.size(0) @@ -50,7 +53,9 @@ def validate(model, loader, device): loss = torch.nn.functional.mse_loss(pred, target) total_loss += loss.item() * 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(): @@ -60,6 +65,13 @@ def main(): 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}") @@ -71,15 +83,27 @@ def main(): cfg.data.data_dir, cfg.data.history_len, cfg.data.action_dim, cfg.data.image_size, cfg.data.normalize_actions, 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( cfg.data.data_dir, cfg.data.history_len, cfg.data.action_dim, cfg.data.image_size, cfg.data.normalize_actions, cfg.data.action_stats_path, val_indices) 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, - 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, shuffle=False, num_workers=cfg.train.num_workers) @@ -97,7 +121,8 @@ def main(): for epoch in range(1, cfg.train.epochs + 1): 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) scheduler.step() print(f" train_loss={train_loss:.6f} val_loss={val_loss:.6f}") @@ -125,13 +150,16 @@ def main(): break 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 action_stats is not None: - stats_path = os.path.join(cfg.train.checkpoint_dir, "action_stats.npy") - np.save(stats_path, action_stats) - print(f"Action stats saved to: {stats_path}") + if best_val_loss == float("inf"): + # 验证集为空时无 best 可选，保底保存最终权重 + torch.save({ + "model_state_dict": model.state_dict(), + "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: with open(args.save_config, "w") as f: