From d17677badedf35c54f1c7b3bce675906c27fc32a Mon Sep 17 00:00:00 2001 From: qing <1776177662@qq.com> Date: Mon, 10 Aug 2026 04:07:37 +0000 Subject: [PATCH] =?UTF-8?q?fix(actpred):=20ONNX=20opset=20=E5=8D=87?= =?UTF-8?q?=E7=BA=A7=20+=20NumPy=20=E5=85=BC=E5=AE=B9=20+=20=E6=96=AD?= =?UTF-8?q?=E7=82=B9=E7=BB=AD=E8=AE=AD=20+=20MobileNetV2=20=E7=A6=BB?= =?UTF-8?q?=E7=BA=BF=E5=BC=B9=E6=80=A7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - export_onnx: opset 17→18，dynamic_axes 改 dynamic_shapes 消除弃用警告 - predict: action_stats 改为 torch tensor，消除 NumPy 2.0 __array_wrap__ 警告 - train/pretrain: 每 epoch 保存 last.pt，新增 --resume 断点续训练 - model: MobileNetEncoder 下载失败时回退到已缓存权重 - _set.sh: 安装时预缓存 MobileNetV2 权重，失败给出明确指引 - test: 新增 TestResumeCheckpoint（2 用例），总计 58 测试全绿 --- actpred/README.md | 10 +++++++ actpred/_set.sh | 19 ++++++++++++ actpred/config.py | 2 ++ actpred/export_onnx.py | 11 +++---- actpred/model.py | 31 +++++++++++++++----- actpred/predict.py | 8 ++--- actpred/pretrain_mt.py | 29 +++++++++++++++++- actpred/test_actpred.py | 65 +++++++++++++++++++++++++++++++++++++++++ actpred/train.py | 26 ++++++++++++++++- 9 files changed, 183 insertions(+), 18 deletions(-) diff --git a/actpred/README.md b/actpred/README.md index 0db7cba..318e43d 100644 --- a/actpred/README.md +++ b/actpred/README.md @@ -30,6 +30,9 @@ cd actpred && bash _set.sh # 训练 bash _run.sh train --data-dir ./data --epochs 100 --action-dim 18 +# 断点续训练（从 last.pt 恢复 model/optimizer/scheduler/epoch） +bash _run.sh train --data-dir ./data --epochs 100 --resume checkpoints/last.pt + # 推理（--images 按时间顺序旧→新传入 frame_len 个图像路径，冷启动时可重复传同一帧； # --action-history 格式：冒号分隔各步，每步 action_dim 个值，共 history_len 步） bash _run.sh --model checkpoints/best_inference.pt \ @@ -55,6 +58,10 @@ python test_actpred.py # 预训练（解析 ../bot4cj/scripts/*.mt） bash _run.sh pretrain --mt-dir ../bot4cj/scripts --epochs 200 --noise-std 1.0 +# 预训练断点续训 +bash _run.sh pretrain --mt-dir ../bot4cj/scripts --epochs 200 --noise-std 1.0 \ + --resume checkpoints/pretrain_last.pt + # 用预训练权重作为微调起点 # 注意：预训练在非归一化动作空间进行，微调时请勿开启 --normalize-actions， # 否则输入分布与预训练权重不匹配，预训练收益会大打折扣 @@ -123,6 +130,9 @@ recorder.flush() | `--seed` | 42 | 随机种子 | | `--num-workers` | 2 | DataLoader 工作进程数 | | `--checkpoint-dir` | ./checkpoints | 检查点保存目录 | +| `--init-from` | 无 | 预训练 checkpoint 路径，作为微调起点（strict=False） | +| `--resume` | 无 | 从 last.pt 断点续训练（恢复 model/optimizer/scheduler/epoch） | +| `--save-config` | 无 | 保存配置到 JSON 文件 | ## 部署注意 diff --git a/actpred/_set.sh b/actpred/_set.sh index f7a315a..3b2d3b8 100755 --- a/actpred/_set.sh +++ b/actpred/_set.sh @@ -15,3 +15,22 @@ pip install onnxscript onnxruntime -i https://pypi.tuna.tsinghua.edu.cn/simple pip check mkdir -p checkpoints data + +# Pre-download MobileNetV2 weights for offline fine-tuning (optional, skip if backbone=simple) +python3 -c " +import torch, sys, os +hub_dir = torch.hub.get_dir() +cache = os.path.join(hub_dir, 'checkpoints', 'mobilenet_v2-b0353104.pth') +if os.path.exists(cache): + print(f'MobileNetV2 weights already cached: {cache}') + sys.exit(0) +url = 'https://download.pytorch.org/models/mobilenet_v2-b0353104.pth' +try: + torch.hub.load_state_dict_from_url(url, model_dir=os.path.join(hub_dir, 'checkpoints')) + print('MobileNetV2 weights cached to: ' + cache) +except Exception as e: + print(f'WARNING: MobileNetV2 weights download failed: {e}') + print(' backbone=simple still works; for mobilenet, manually download from:') + print(f' {url}') + print(f' and place at: {cache}') +" || echo "WARNING: MobileNetV2 weight pre-download skipped" diff --git a/actpred/config.py b/actpred/config.py index 395aa54..21c4dd7 100644 --- a/actpred/config.py +++ b/actpred/config.py @@ -85,6 +85,8 @@ def build_argparser(description: str = "") -> argparse.ArgumentParser: parser.add_argument("--max-grad-norm", type=float, default=None) parser.add_argument("--init-from", type=str, default=None, help="预训练 checkpoint 路径，作为微调起点（strict=False 加载兼容权重）") + parser.add_argument("--resume", type=str, default=None, + help="从 last.pt 断点续训（恢复 model/optimizer/scheduler/epoch）") parser.add_argument("--save-config", type=str, default=None, help="保存配置到 JSON 文件") return parser diff --git a/actpred/export_onnx.py b/actpred/export_onnx.py index 06befac..55ab562 100644 --- a/actpred/export_onnx.py +++ b/actpred/export_onnx.py @@ -1,5 +1,6 @@ import argparse import torch +from torch.export import Dim from config import parse_checkpoint_config from model import ActionPredictor @@ -23,17 +24,17 @@ def main(): 1, data_cfg.frame_len, 3, data_cfg.image_size, data_cfg.image_size) dummy_history = torch.randn(1, data_cfg.history_len, data_cfg.action_dim) + batch = Dim("batch") torch.onnx.export( model, (dummy_images, dummy_history), args.output, input_names=["images", "action_history"], output_names=["predicted_action"], - dynamic_axes={ - "images": {0: "batch"}, - "action_history": {0: "batch"}, - "predicted_action": {0: "batch"}, + dynamic_shapes={ + "images": {0: batch}, + "action_history": {0: batch}, }, - opset_version=17, + opset_version=18, ) print(f"ONNX model exported to: {args.output}") diff --git a/actpred/model.py b/actpred/model.py index 7a6d6fc..9eaa817 100644 --- a/actpred/model.py +++ b/actpred/model.py @@ -36,13 +36,30 @@ class MobileNetEncoder(nn.Module): except Exception as e: if not pretrained: raise RuntimeError("加载 MobileNetV2 架构失败") from e - raise RuntimeError( - "MobileNetV2 预训练权重下载失败，可能是网络受限。" - "请在联网环境中预下载权重" - "（运行 torchvision.models.mobilenet_v2(weights='IMAGENET1K_V1')），" - "并将 ~/.cache/torch/hub/checkpoints/ 下的权重文件" - "拷贝到目标机器同路径。" - ) from e + # 主源下载失败，尝试加载已缓存的权重文件 + try: + import os + hub_dir = torch.hub.get_dir() + cache_dir = os.path.join(hub_dir, "checkpoints") + weight_file = getattr(weights, "file_name", + "mobilenet_v2-b0353104.pth") + cached_path = os.path.join(cache_dir, weight_file) + if os.path.exists(cached_path): + print(f"INFO: 使用缓存权重 {cached_path}") + backbone = mobilenet_v2(weights=None) + backbone.load_state_dict(torch.load( + cached_path, map_location="cpu", + weights_only=True)) + else: + raise FileNotFoundError(cached_path) + except Exception: + raise RuntimeError( + "MobileNetV2 预训练权重下载失败，可能是网络受限。" + "请在联网环境中预下载权重" + "（运行 torchvision.models.mobilenet_v2(weights='IMAGENET1K_V1')），" + "并将 ~/.cache/torch/hub/checkpoints/ 下的权重文件" + "拷贝到目标机器同路径。" + ) from e self.features = backbone.features self.pool = nn.AdaptiveAvgPool2d(1) self.proj = nn.Linear(1280, feature_dim) diff --git a/actpred/predict.py b/actpred/predict.py index c9eb971..f09cf3c 100644 --- a/actpred/predict.py +++ b/actpred/predict.py @@ -58,8 +58,8 @@ class ActionPredictorInference: raise ValueError( f"action_stats.npy 形状应为 (2, {self._data_cfg.action_dim})，" f"实际为 {stats.shape}，请检查文件: {stats_path}") - self._action_mean = stats[0].astype(np.float32) - self._action_std = stats[1].astype(np.float32) + self._action_mean = torch.from_numpy(stats[0].astype(np.float32)).to(self._device) + self._action_std = torch.from_numpy(stats[1].astype(np.float32)).to(self._device) @torch.no_grad() def predict(self, images, action_history: np.ndarray) -> np.ndarray: @@ -83,10 +83,10 @@ class ActionPredictorInference: if self._action_mean is not None: history = (history - self._action_mean) / self._action_std history = history.unsqueeze(0).to(self._device) - pred = self._model(img_tensor, history).squeeze(0).cpu().numpy() + pred = self._model(img_tensor, history).squeeze(0) if self._action_mean is not None: pred = pred * self._action_std + self._action_mean - return pred + return pred.cpu().numpy() def parse_action_history(s: str, action_dim: int) -> np.ndarray: diff --git a/actpred/pretrain_mt.py b/actpred/pretrain_mt.py index ed18375..e3b06b4 100644 --- a/actpred/pretrain_mt.py +++ b/actpred/pretrain_mt.py @@ -216,12 +216,25 @@ def main(): best_val_loss = float("inf") patience_counter = 0 + start_epoch = 1 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" ) + last_path = os.path.join(cfg.train.checkpoint_dir, "pretrain_last.pt") - for epoch in range(1, cfg.train.epochs + 1): + if args.resume: + resume_ckpt = torch.load(args.resume, map_location=device, weights_only=True) + model.load_state_dict(resume_ckpt["model_state_dict"]) + optimizer.load_state_dict(resume_ckpt["optimizer_state_dict"]) + scheduler.load_state_dict(resume_ckpt["scheduler_state_dict"]) + start_epoch = resume_ckpt["epoch"] + 1 + best_val_loss = resume_ckpt.get("best_val_loss", float("inf")) + patience_counter = resume_ckpt.get("patience_counter", 0) + print(f"Resumed from {args.resume} at epoch {start_epoch} " + f"(best_val_loss={best_val_loss:.6f})") + + for epoch in range(start_epoch, cfg.train.epochs + 1): print(f"\nEpoch {epoch}/{cfg.train.epochs} " f"lr={optimizer.param_groups[0]['lr']:.6f}") train_loss = pretrain_one_epoch( @@ -262,6 +275,20 @@ def main(): f"(patience={cfg.train.patience})") break + # 每 epoch 保存 last.pt，支持 --resume 断点续训练 + torch.save( + { + "epoch": epoch, + "model_state_dict": model.state_dict(), + "optimizer_state_dict": optimizer.state_dict(), + "scheduler_state_dict": scheduler.state_dict(), + "best_val_loss": best_val_loss, + "patience_counter": patience_counter, + "config": cfg.to_dict(), + }, + last_path, + ) + print(f"\nPre-training done. Best val_loss={best_val_loss:.6f}") if best_val_loss == float("inf"): diff --git a/actpred/test_actpred.py b/actpred/test_actpred.py index a92a6ad..84e426e 100644 --- a/actpred/test_actpred.py +++ b/actpred/test_actpred.py @@ -717,5 +717,70 @@ class TestExportOnnx(unittest.TestCase): self.assertGreater(os.path.getsize(out_path), 0) +class TestResumeCheckpoint(unittest.TestCase): + """测试 --resume 断点续训练：checkpoint 格式与状态恢复""" + + def test_last_checkpoint_contains_resume_keys(self): + """last.pt 必须包含恢复训练所需的全部键""" + required_keys = { + "epoch", "model_state_dict", "optimizer_state_dict", + "scheduler_state_dict", "best_val_loss", "patience_counter", + "config", + } + cfg = Config() + cfg.data.action_dim = 4 + model = ActionPredictor(cfg.model, cfg.data) + optimizer = torch.optim.AdamW(model.parameters(), lr=0.001) + scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( + optimizer, T_max=100) + with tempfile.TemporaryDirectory() as tmpdir: + last_path = os.path.join(tmpdir, "last.pt") + torch.save({ + "epoch": 10, + "model_state_dict": model.state_dict(), + "optimizer_state_dict": optimizer.state_dict(), + "scheduler_state_dict": scheduler.state_dict(), + "best_val_loss": 0.5, + "patience_counter": 3, + "config": cfg.to_dict(), + }, last_path) + ckpt = torch.load(last_path, map_location="cpu", weights_only=True) + self.assertTrue(required_keys.issubset(ckpt.keys()), + f"缺少键: {required_keys - set(ckpt.keys())}") + self.assertEqual(ckpt["epoch"], 10) + self.assertEqual(ckpt["best_val_loss"], 0.5) + self.assertEqual(ckpt["patience_counter"], 3) + + def test_resume_restores_model_weights(self): + """resume 加载后模型权重应与保存时一致""" + cfg = Config() + cfg.data.action_dim = 4 + torch.manual_seed(0) + model = ActionPredictor(cfg.model, cfg.data) + optimizer = torch.optim.AdamW(model.parameters(), lr=0.001) + scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( + optimizer, T_max=100) + with tempfile.TemporaryDirectory() as tmpdir: + last_path = os.path.join(tmpdir, "last.pt") + torch.save({ + "epoch": 5, + "model_state_dict": model.state_dict(), + "optimizer_state_dict": optimizer.state_dict(), + "scheduler_state_dict": scheduler.state_dict(), + "best_val_loss": 1.0, + "patience_counter": 0, + "config": cfg.to_dict(), + }, last_path) + # 重新创建模型并恢复 + model2 = ActionPredictor(cfg.model, cfg.data) + ckpt = torch.load(last_path, map_location="cpu", weights_only=True) + model2.load_state_dict(ckpt["model_state_dict"]) + for k, v in model.state_dict().items(): + self.assertTrue(torch.equal(model2.state_dict()[k], v), + msg=f"权重不一致: {k}") + + + + if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/actpred/train.py b/actpred/train.py index 50b8ffe..90ced88 100644 --- a/actpred/train.py +++ b/actpred/train.py @@ -132,9 +132,22 @@ def main(): best_val_loss = float("inf") patience_counter = 0 + start_epoch = 1 best_path = os.path.join(cfg.train.checkpoint_dir, "best.pt") + last_path = os.path.join(cfg.train.checkpoint_dir, "last.pt") - for epoch in range(1, cfg.train.epochs + 1): + if args.resume: + resume_ckpt = torch.load(args.resume, map_location=device, weights_only=True) + model.load_state_dict(resume_ckpt["model_state_dict"]) + optimizer.load_state_dict(resume_ckpt["optimizer_state_dict"]) + scheduler.load_state_dict(resume_ckpt["scheduler_state_dict"]) + start_epoch = resume_ckpt["epoch"] + 1 + best_val_loss = resume_ckpt.get("best_val_loss", float("inf")) + patience_counter = resume_ckpt.get("patience_counter", 0) + print(f"Resumed from {args.resume} at epoch {start_epoch} " + f"(best_val_loss={best_val_loss:.6f})") + + for epoch in range(start_epoch, 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, cfg.train.max_grad_norm) @@ -164,6 +177,17 @@ def main(): print(f"\nEarly stopping at epoch {epoch} (patience={cfg.train.patience})") break + # 每 epoch 保存 last.pt，支持 --resume 断点续训 + torch.save({ + "epoch": epoch, + "model_state_dict": model.state_dict(), + "optimizer_state_dict": optimizer.state_dict(), + "scheduler_state_dict": scheduler.state_dict(), + "best_val_loss": best_val_loss, + "patience_counter": patience_counter, + "config": cfg.to_dict(), + }, last_path) + print(f"\nTraining done. Best val_loss={best_val_loss:.6f}") if best_val_loss == float("inf"):