fix(actpred): ONNX opset 升级 + NumPy 兼容 + 断点续训 + MobileNetV2 离线弹性

- 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

@@ -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 文件 |
 
 ## 部署注意
 

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"

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
 

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={
+        dynamic_shapes={
-            "images": {0: "batch"},
+            "images": {0: batch},
-            "action_history": {0: "batch"},
+            "action_history": {0: batch},
-            "predicted_action": {0: "batch"},
         },
-        opset_version=17,
+        opset_version=18,
     )
     print(f"ONNX model exported to: {args.output}")
 

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 预训练权重下载失败，可能是网络受限。"
+            try:
-                "请在联网环境中预下载权重"
+                import os
-                "（运行 torchvision.models.mobilenet_v2(weights='IMAGENET1K_V1')），"
+                hub_dir = torch.hub.get_dir()
-                "并将 ~/.cache/torch/hub/checkpoints/ 下的权重文件"
+                cache_dir = os.path.join(hub_dir, "checkpoints")
-                "拷贝到目标机器同路径。"
+                weight_file = getattr(weights, "file_name",
-            ) from e
+                                      "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)

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_mean = torch.from_numpy(stats[0].astype(np.float32)).to(self._device)
-            self._action_std = stats[1].astype(np.float32)
+            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:

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"):

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)

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"):