From 4dd25486f830eb5330fa015a4007a6e0558ca7d2 Mon Sep 17 00:00:00 2001 From: qing <1776177662@qq.com> Date: Sun, 9 Aug 2026 18:17:03 +0000 Subject: [PATCH] =?UTF-8?q?feat(actpred):=20=E5=A4=9A=E5=B8=A7=E6=B3=A8?= =?UTF-8?q?=E6=84=8F=E5=8A=9B=E6=9E=B6=E6=9E=84=E6=94=B9=E9=80=A0=EF=BC=88?= =?UTF-8?q?Transformer=20=E6=9B=BF=E6=8D=A2=20concat=20=E8=9E=8D=E5=90=88?= =?UTF-8?q?=EF=BC=8C=E6=94=AF=E6=8C=81=E5=A4=9A=E5=B8=A7=E5=9B=BE=E5=83=8F?= =?UTF-8?q?+=E5=8E=86=E5=8F=B2=E5=8A=A8=E4=BD=9C=E5=BA=8F=E5=88=97?= =?UTF-8?q?=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- actpred/README.md | 42 ++++++++---- actpred/config.py | 28 ++++++-- actpred/dataset.py | 22 +++++-- actpred/export_onnx.py | 15 +++-- actpred/model.py | 90 +++++++++++++------------- actpred/predict.py | 29 +++++++-- actpred/pretrain_mt.py | 36 ++++++----- actpred/test_actpred.py | 137 ++++++++++++++++++++++++++++++---------- actpred/train.py | 20 +++--- 9 files changed, 272 insertions(+), 147 deletions(-) diff --git a/actpred/README.md b/actpred/README.md index f46163d..0db7cba 100644 --- a/actpred/README.md +++ b/actpred/README.md @@ -1,19 +1,26 @@ # actpred — 动作预测模块 -输入摄像头图像 + 历史动作序列,预测下一个连续动作(回归)。 +输入多帧摄像头图像 + 历史动作序列,预测下一时间步的连续动作(回归)。 -纯神经网络实现(PyTorch),支持轻量 CNN 或 MobileNetV2 图像编码,MLP/GRU 动作历史编码。总参数量 ~140K(默认)至 ~2.3M。 +纯神经网络实现(PyTorch):轻量 CNN 或 MobileNetV2 逐帧编码,图像/动作 token 进轻量 +Transformer(pre-LN)自注意力融合,取 [PRED] token 回归下一步动作。 +默认配置总参数量 ~450K。 ## 架构 ``` -图像 (3×64×64) ──► CNN 编码器 ──► img_feat (128) - │ - ├─ concat ──► 融合头 ──► 预测动作 (action_dim) - │ -动作历史 (H×D) ──► MLP/GRU ──► act_feat (64) +F 帧图像 (F×3×64×64) ─► 共享 CNN 编码器 ─► F 个 token ─► proj → d_model ┐ + ├─► [img×F, act×H, PRED] + 位置编码 +动作历史 (H×D) ─► Linear proj → d_model ─► H 个 token ────────┤ │ +可学习 [PRED] token ────────────────────────────────────────┘ ▼ + TransformerEncoder (pre-LN, 2 层, 4 头, d_model=128) + │ + 取 PRED 位置输出 ─► MLP 头 ─► 预测动作 (action_dim) ``` +token 序列长度固定为 F+H+1(默认 3+5+1=9),除 batch 外无动态 shape,对 NPU(如树莓派 +AI 扩展板)部署友好。 + ## 快速开始 ```bash @@ -23,8 +30,10 @@ cd actpred && bash _set.sh # 训练 bash _run.sh train --data-dir ./data --epochs 100 --action-dim 18 -# 推理(--action-history 格式:冒号分隔各步,每步 action_dim 个值,共 history_len 步) -bash _run.sh --model checkpoints/best_inference.pt --image frame.jpg \ +# 推理(--images 按时间顺序旧→新传入 frame_len 个图像路径,冷启动时可重复传同一帧; +# --action-history 格式:冒号分隔各步,每步 action_dim 个值,共 history_len 步) +bash _run.sh --model checkpoints/best_inference.pt \ + --images frame_t-2.jpg frame_t-1.jpg frame_t.jpg \ --action-history "45,0,45,0,45,0,45,0,90,90,10,200,0,0,0,0,0,0:45,0,45,0,45,0,45,0,90,90,0,20,0,0,0,0,0,0:45,0,45,0,45,0,45,0,90,90,0,0,0,0,0,0,0,0:45,0,45,0,45,0,45,0,90,90,0,0,0,0,0,0,0,0:45,0,45,0,45,0,45,0,90,90,0,0,0,0,0,0,0,0" # 导出 ONNX(如安装了 onnxruntime,导出后自动验证一致性) @@ -39,7 +48,7 @@ python test_actpred.py ## 预训练(从 .mt 脚本,无需图像) -从 bot4cj 的 `.mt` 动作脚本中提取 `@Action` 矩阵,用零图像占位预训练动作动力学(history → next action), +从 bot4cj 的 `.mt` 动作脚本中提取 `@Action` 矩阵,用全零多帧图像占位预训练动作动力学(history → next action), 保存的 checkpoint 与正常训练格式完全兼容,可直接用于推理或作为微调起点。 ```bash @@ -74,7 +83,8 @@ data/ ... ``` -每个 timestep 的样本:当前图像 + 前 `history_len` 步动作 → 当前步动作。 +每个 timestep 的样本:帧 `[t-frame_len+1, t]`(旧→新)+ 前 `history_len` 步动作 → 当前步动作。 +采样起点为 `max(history_len, frame_len-1)`。 ### 数据采集 @@ -96,8 +106,11 @@ recorder.flush() | `--image-size` | 64 | 图像缩放尺寸 | | `--action-dim` | 18 | 动作向量维度(对齐 bot4cj 18 列:10 舵机 + 2 云台 + 扬声器频率 + 延时 + 4 预留)| | `--history-len` | 5 | 历史动作窗口长度 | +| `--frame-len` | 3 | 输入图像帧数(含当前帧,旧帧在前)| | `--backbone` | simple | CNN 类型 (simple/mobilenet) | -| `--history-encoder` | mlp | 历史编码器 (mlp/gru) | +| `--attn-dim` | 128 | 注意力 token 维度 d_model(须能被 attn-heads 整除)| +| `--attn-heads` | 4 | 注意力头数 | +| `--attn-layers` | 2 | TransformerEncoder 层数 | | `--epochs` | 100 | 训练轮数 | | `--lr` | 0.001 | 学习率 | | `--batch-size` | 32 | 批次大小 | @@ -124,7 +137,8 @@ recorder.flush() from predict import ActionPredictorInference predictor = ActionPredictorInference("checkpoints/best_inference.pt") -print(predictor.action_dim, predictor.history_len) # 公共属性 -action = predictor.predict(image_rgb, action_history) # np.ndarray (action_dim,) +print(predictor.action_dim, predictor.history_len, predictor.frame_len) # 公共属性 +# images: 长度为 frame_len 的 RGB 图像列表(旧→新,末帧为当前帧;冷启动可重复同一帧) +action = predictor.predict(images, action_history) # np.ndarray (action_dim,) # action_history 形状须为 (history_len, action_dim) ``` diff --git a/actpred/config.py b/actpred/config.py index f6e1e98..395aa54 100644 --- a/actpred/config.py +++ b/actpred/config.py @@ -8,6 +8,7 @@ class DataConfig: image_size: int = 64 action_dim: int = 18 history_len: int = 5 + frame_len: int = 3 data_dir: str = "./data" normalize_actions: bool = False action_stats_path: Optional[str] = None @@ -17,10 +18,11 @@ class DataConfig: class ModelConfig: cnn_backbone: str = "simple" cnn_feature_dim: int = 128 - action_hidden_dim: int = 64 fused_hidden_dim: int = 128 dropout: float = 0.1 - history_encoder: str = "mlp" + attn_dim: int = 128 + attn_heads: int = 4 + attn_layers: int = 2 @dataclass @@ -60,15 +62,17 @@ def build_argparser(description: str = "") -> argparse.ArgumentParser: parser.add_argument("--image-size", type=int, default=None) parser.add_argument("--action-dim", type=int, default=None) parser.add_argument("--history-len", type=int, default=None) + parser.add_argument("--frame-len", type=int, default=None) parser.add_argument("--data-dir", type=str, default=None) parser.add_argument("--normalize-actions", action=argparse.BooleanOptionalAction, default=None) parser.add_argument("--action-stats-path", type=str, default=None) parser.add_argument("--backbone", type=str, default=None, choices=["simple", "mobilenet"]) parser.add_argument("--cnn-feature-dim", type=int, default=None) - parser.add_argument("--action-hidden-dim", type=int, default=None) parser.add_argument("--fused-hidden-dim", type=int, default=None) parser.add_argument("--dropout", type=float, default=None) - parser.add_argument("--history-encoder", type=str, default=None, choices=["mlp", "gru"]) + parser.add_argument("--attn-dim", type=int, default=None) + parser.add_argument("--attn-heads", type=int, default=None) + parser.add_argument("--attn-layers", type=int, default=None) parser.add_argument("--batch-size", type=int, default=None) parser.add_argument("--epochs", type=int, default=None) parser.add_argument("--lr", type=float, default=None) @@ -95,11 +99,14 @@ def validate_config_values(data_cfg: DataConfig, model_cfg: ModelConfig, checks = [ ("image_size", data_cfg.image_size, 1, 1024, int), ("history_len", data_cfg.history_len, 1, 1024, int), + ("frame_len", data_cfg.frame_len, 1, 1024, int), ("action_dim", data_cfg.action_dim, 1, 1024, int), ("cnn_feature_dim", model_cfg.cnn_feature_dim, 1, 4096, int), - ("action_hidden_dim", model_cfg.action_hidden_dim, 1, 4096, int), ("fused_hidden_dim", model_cfg.fused_hidden_dim, 1, 4096, int), ("dropout", model_cfg.dropout, 0.0, 1.0, float), + ("attn_dim", model_cfg.attn_dim, 1, 4096, int), + ("attn_heads", model_cfg.attn_heads, 1, 64, int), + ("attn_layers", model_cfg.attn_layers, 1, 64, int), ] for name, val, lo, hi, typ in checks: if typ is int: @@ -110,6 +117,11 @@ def validate_config_values(data_cfg: DataConfig, model_cfg: ModelConfig, raise ValueError( f"checkpoint config {name}={val} 超出合理范围 [{lo}, {hi}]") + if model_cfg.attn_dim % model_cfg.attn_heads != 0: + raise ValueError( + f"attn_dim={model_cfg.attn_dim} 必须能被 attn_heads=" + f"{model_cfg.attn_heads} 整除") + if train_cfg is not None: t_checks = [ ("batch_size", train_cfg.batch_size, 1, 65536, int), @@ -162,15 +174,17 @@ def apply_overrides(cfg: Config, args: argparse.Namespace) -> Config: "image_size": ("data", "image_size"), "action_dim": ("data", "action_dim"), "history_len": ("data", "history_len"), + "frame_len": ("data", "frame_len"), "data_dir": ("data", "data_dir"), "normalize_actions": ("data", "normalize_actions"), "action_stats_path": ("data", "action_stats_path"), "backbone": ("model", "cnn_backbone"), "cnn_feature_dim": ("model", "cnn_feature_dim"), - "action_hidden_dim": ("model", "action_hidden_dim"), "fused_hidden_dim": ("model", "fused_hidden_dim"), "dropout": ("model", "dropout"), - "history_encoder": ("model", "history_encoder"), + "attn_dim": ("model", "attn_dim"), + "attn_heads": ("model", "attn_heads"), + "attn_layers": ("model", "attn_layers"), "batch_size": ("train", "batch_size"), "epochs": ("train", "epochs"), "lr": ("train", "lr"), diff --git a/actpred/dataset.py b/actpred/dataset.py index 54a14ae..fd4e924 100644 --- a/actpred/dataset.py +++ b/actpred/dataset.py @@ -17,10 +17,12 @@ def _natural_key(p): class ActionDataset(Dataset): def __init__(self, data_dir: str, history_len: int, action_dim: int, image_size: int = 64, normalize_actions: bool = False, - action_stats_path: str = None, episode_indices: list = None): + action_stats_path: str = None, episode_indices: list = None, + frame_len: int = 3): self.history_len = history_len self.action_dim = action_dim self.image_size = image_size + self.frame_len = frame_len self.transform = transforms.Compose([ transforms.Resize((image_size, image_size)), @@ -51,8 +53,11 @@ class ActionDataset(Dataset): print(f"WARNING: {ep_dir} 动作帧数 {n_frames} 与图像数 {n_images} 不一致," f"按较少者({min(n_frames, n_images)})使用,请检查录制完整性") n = min(n_frames, n_images) - for t in range(history_len, n): - self.samples.append((image_files[t], actions[t - history_len:t], actions[t])) + # 每个样本:帧 [t-frame_len+1, t](旧→新)+ 动作历史 [t-history_len, t) → 动作 t + start = max(history_len, frame_len - 1) + for t in range(start, n): + self.samples.append((image_files[t - frame_len + 1: t + 1], + actions[t - history_len:t], actions[t])) self._action_mean = None self._action_std = None @@ -77,15 +82,18 @@ class ActionDataset(Dataset): return len(self.samples) def __getitem__(self, idx): - image_path, action_history, target = self.samples[idx] - image = Image.open(image_path).convert("RGB") - image = self.transform(image) + image_paths, action_history, target = self.samples[idx] + frames = [] + for p in image_paths: + image = Image.open(p).convert("RGB") + frames.append(self.transform(image)) + images = torch.stack(frames) # (F, 3, S, S) action_history = torch.tensor(action_history, dtype=torch.float32) target = torch.tensor(target, dtype=torch.float32) if self._action_mean is not None: action_history = (action_history - self._action_mean) / self._action_std target = (target - self._action_mean) / self._action_std - return {"image": image, "action_history": action_history, "target": target} + return {"images": images, "action_history": action_history, "target": target} def get_action_stats(self): if self._action_mean is not None: diff --git a/actpred/export_onnx.py b/actpred/export_onnx.py index ac367fd..06befac 100644 --- a/actpred/export_onnx.py +++ b/actpred/export_onnx.py @@ -19,16 +19,17 @@ def main(): model.load_state_dict(checkpoint["model_state_dict"]) model.eval() - dummy_image = torch.randn(1, 3, data_cfg.image_size, data_cfg.image_size) + dummy_images = torch.randn( + 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) torch.onnx.export( - model, (dummy_image, dummy_history), + model, (dummy_images, dummy_history), args.output, - input_names=["image", "action_history"], + input_names=["images", "action_history"], output_names=["predicted_action"], dynamic_axes={ - "image": {0: "batch"}, + "images": {0: "batch"}, "action_history": {0: "batch"}, "predicted_action": {0: "batch"}, }, @@ -41,11 +42,11 @@ def main(): import onnxruntime as ort import numpy as np sess = ort.InferenceSession(args.output) - img_np = dummy_image.numpy() + img_np = dummy_images.numpy() hist_np = dummy_history.numpy() - ort_out = sess.run(None, {"image": img_np, "action_history": hist_np})[0] + ort_out = sess.run(None, {"images": img_np, "action_history": hist_np})[0] with torch.no_grad(): - torch_out = model(dummy_image, dummy_history).numpy() + torch_out = model(dummy_images, dummy_history).numpy() max_diff = np.abs(ort_out - torch_out).max() print(f"ONNX 一致性验证通过(最大差异: {max_diff:.6e})") except ImportError: diff --git a/actpred/model.py b/actpred/model.py index 30378c4..7a6d6fc 100644 --- a/actpred/model.py +++ b/actpred/model.py @@ -53,34 +53,7 @@ class MobileNetEncoder(nn.Module): return self.proj(x) -class MLPHistoryEncoder(nn.Module): - def __init__(self, input_dim: int, hidden_dim: int = 64, dropout: float = 0.1): - super().__init__() - self.net = nn.Sequential( - nn.Linear(input_dim, hidden_dim), - nn.ReLU(inplace=True), - nn.Dropout(dropout), - nn.Linear(hidden_dim, hidden_dim), - nn.ReLU(inplace=True), - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.net(x) - - -class GRUHistoryEncoder(nn.Module): - def __init__(self, input_dim: int, hidden_dim: int = 64, dropout: float = 0.1): - super().__init__() - self.gru = nn.GRU(input_size=input_dim, hidden_size=hidden_dim, - num_layers=1, batch_first=True) - self.dropout = nn.Dropout(dropout) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - _, h = self.gru(x) - return self.dropout(h.squeeze(0)) - - -class FusionHead(nn.Module): +class PredictionHead(nn.Module): def __init__(self, input_dim: int, hidden_dim: int, output_dim: int, dropout: float = 0.1): super().__init__() self.net = nn.Sequential( @@ -97,6 +70,13 @@ class FusionHead(nn.Module): class ActionPredictor(nn.Module): + """多帧图像 + 历史动作的 Transformer 注意力融合预测器 + + token 序列:[img×F, act×H, PRED],长度固定为 F+H+1, + 取 PRED 位置输出进预测头回归下一时间步动作。 + 除 batch 维外全部维度静态,保证 ONNX/NPU 部署友好。 + """ + def __init__(self, model_cfg: ModelConfig, data_cfg: DataConfig, pretrained: bool = True): super().__init__() @@ -106,26 +86,44 @@ class ActionPredictor(nn.Module): else: self.image_encoder = SimpleCNN(model_cfg.cnn_feature_dim) - history_input_dim = data_cfg.history_len * data_cfg.action_dim - if model_cfg.history_encoder == "gru": - self.action_encoder = GRUHistoryEncoder( - data_cfg.action_dim, model_cfg.action_hidden_dim, model_cfg.dropout) - else: - self.action_encoder = MLPHistoryEncoder( - history_input_dim, model_cfg.action_hidden_dim, model_cfg.dropout) + d = model_cfg.attn_dim + self.img_proj = nn.Linear(model_cfg.cnn_feature_dim, d) + self.act_proj = nn.Linear(data_cfg.action_dim, d) - fused_input = model_cfg.cnn_feature_dim + model_cfg.action_hidden_dim - self.head = FusionHead( - fused_input, model_cfg.fused_hidden_dim, data_cfg.action_dim, model_cfg.dropout) + seq_len = data_cfg.frame_len + data_cfg.history_len + 1 + self.pred_token = nn.Parameter(torch.zeros(1, 1, d)) + self.pos_embed = nn.Parameter(torch.zeros(1, seq_len, d)) + nn.init.trunc_normal_(self.pred_token, std=0.02) + nn.init.trunc_normal_(self.pos_embed, std=0.02) - self._use_gru = model_cfg.history_encoder == "gru" + encoder_layer = nn.TransformerEncoderLayer( + d_model=d, nhead=model_cfg.attn_heads, + dim_feedforward=d * 2, dropout=model_cfg.dropout, + batch_first=True, norm_first=True) + self.encoder = nn.TransformerEncoder( + encoder_layer, num_layers=model_cfg.attn_layers, + enable_nested_tensor=False) + + self.head = PredictionHead( + d, model_cfg.fused_hidden_dim, data_cfg.action_dim, model_cfg.dropout) + + self._frame_len = data_cfg.frame_len self._history_len = data_cfg.history_len self._action_dim = data_cfg.action_dim + self._image_size = data_cfg.image_size + self._cnn_feature_dim = model_cfg.cnn_feature_dim - def forward(self, image: torch.Tensor, action_history: torch.Tensor) -> torch.Tensor: - img_feat = self.image_encoder(image) - if self._use_gru: - act_feat = self.action_encoder(action_history) - else: - act_feat = self.action_encoder(action_history.flatten(1)) - return self.head(torch.cat([img_feat, act_feat], dim=-1)) + def forward(self, images: torch.Tensor, action_history: torch.Tensor) -> torch.Tensor: + # images: (B, F, 3, S, S) action_history: (B, H, D) + b = images.size(0) + frames = images.reshape( + b * self._frame_len, 3, self._image_size, self._image_size) + img_feat = self.image_encoder(frames).reshape( + b, self._frame_len, self._cnn_feature_dim) + img_tokens = self.img_proj(img_feat) + act_tokens = self.act_proj(action_history) + pred_token = self.pred_token.expand(b, -1, -1) + tokens = torch.cat([img_tokens, act_tokens, pred_token], dim=1) + tokens = tokens + self.pos_embed + out = self.encoder(tokens) + return self.head(out[:, -1]) diff --git a/actpred/predict.py b/actpred/predict.py index 165681d..c9eb971 100644 --- a/actpred/predict.py +++ b/actpred/predict.py @@ -28,6 +28,7 @@ class ActionPredictorInference: self.action_dim = self._data_cfg.action_dim self.history_len = self._data_cfg.history_len + self.frame_len = self._data_cfg.frame_len self._transform = transforms.Compose([ transforms.Resize((self._data_cfg.image_size, self._data_cfg.image_size)), @@ -61,13 +62,23 @@ class ActionPredictorInference: self._action_std = stats[1].astype(np.float32) @torch.no_grad() - def predict(self, image: np.ndarray, action_history: np.ndarray) -> np.ndarray: + def predict(self, images, action_history: np.ndarray) -> np.ndarray: + """images: 长度为 frame_len 的 RGB 图像序列(旧→新,末帧为当前帧)""" + if not isinstance(images, (list, tuple)): + images = [images] + if len(images) != self._data_cfg.frame_len: + raise ValueError( + f"images 应为 {self._data_cfg.frame_len} 帧,实际为 {len(images)} 帧。" + f"冷启动时可重复传入同一帧。") expected_shape = (self._data_cfg.history_len, self._data_cfg.action_dim) if action_history.shape != expected_shape: raise ValueError( f"action_history 形状应为 {expected_shape},实际为 {action_history.shape}") - img = Image.fromarray(image) if isinstance(image, np.ndarray) else image - img_tensor = self._transform(img).unsqueeze(0).to(self._device) + frames = [] + for image in images: + img = Image.fromarray(image) if isinstance(image, np.ndarray) else image + frames.append(self._transform(img)) + img_tensor = torch.stack(frames).unsqueeze(0).to(self._device) # (1, F, 3, S, S) history = torch.tensor(action_history, dtype=torch.float32) if self._action_mean is not None: history = (history - self._action_mean) / self._action_std @@ -93,7 +104,9 @@ def parse_action_history(s: str, action_dim: int) -> np.ndarray: def main(): parser = argparse.ArgumentParser(description="actpred inference") parser.add_argument("--model", required=True, help="checkpoint path") - parser.add_argument("--image", required=True, help="image file path") + parser.add_argument("--images", required=True, nargs="+", + help="按时间顺序(旧→新)传入 frame_len 个图像路径;" + "冷启动时可重复传同一帧路径") parser.add_argument("--action-history", required=True, help="冒号分隔各步,每步 action_dim 个值,共 history_len 步。" "e.g. '0,0,0,0:1,0,0,0' (action_dim=4, history_len=2)") @@ -106,9 +119,13 @@ def main(): predictor = ActionPredictorInference(args.model, args.device) action_dim = args.action_dim or predictor.action_dim - image = Image.open(args.image).convert("RGB") + if len(args.images) != predictor.frame_len: + raise SystemExit( + f"--images 需传入 {predictor.frame_len} 个路径(旧→新)," + f"实际为 {len(args.images)} 个。冷启动时可重复传同一帧路径。") + images = [np.array(Image.open(p).convert("RGB")) for p in args.images] history = parse_action_history(args.action_history, action_dim) - predicted = predictor.predict(np.array(image), history) + predicted = predictor.predict(images, history) if args.output_format == "json": print(json.dumps({"predicted_action": predicted.tolist()}, indent=2)) diff --git a/actpred/pretrain_mt.py b/actpred/pretrain_mt.py index dbeb29e..ed18375 100644 --- a/actpred/pretrain_mt.py +++ b/actpred/pretrain_mt.py @@ -76,24 +76,25 @@ class MTActionDataset(Dataset): } -def pretrain_one_epoch(model, loader, optimizer, device, image_size, +def pretrain_one_epoch(model, loader, optimizer, device, image_size, frame_len, max_grad_norm=None, noise_std=0.0): model.train() total_loss = 0.0 count = 0 - zero_image = None + zero_images = 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 + if zero_images is None or zero_images.size(0) != history.size(0): + zero_images = torch.zeros( + history.size(0), frame_len, 3, image_size, image_size, + device=device, ) optimizer.zero_grad() - pred = model(zero_image, history) + pred = model(zero_images, history) loss = torch.nn.functional.mse_loss(pred, target) loss.backward() if max_grad_norm: @@ -105,19 +106,20 @@ def pretrain_one_epoch(model, loader, optimizer, device, image_size, @torch.no_grad() -def validate(model, loader, device, image_size): +def validate(model, loader, device, image_size, frame_len): model.eval() total_loss = 0.0 count = 0 - zero_image = None + zero_images = 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 + if zero_images is None or zero_images.size(0) != history.size(0): + zero_images = torch.zeros( + history.size(0), frame_len, 3, image_size, image_size, + device=device, ) - pred = model(zero_image, history) + pred = model(zero_images, history) loss = torch.nn.functional.mse_loss(pred, target) total_loss += loss.item() * history.size(0) count += history.size(0) @@ -198,8 +200,10 @@ def main(): print(f"Model params: {n_params:,}") print(f"Config: action_dim={cfg.data.action_dim}, " f"history_len={cfg.data.history_len}, " + f"frame_len={cfg.data.frame_len}, " f"backbone={cfg.model.cnn_backbone}, " - f"encoder={cfg.model.history_encoder}") + f"attn={cfg.model.attn_layers}x{cfg.model.attn_heads}h" + f"@{cfg.model.attn_dim}") print(f"Pre-training: noise_std={args.noise_std}") optimizer = torch.optim.AdamW( @@ -222,9 +226,11 @@ def main(): 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, + cfg.data.image_size, cfg.data.frame_len, + cfg.train.max_grad_norm, args.noise_std, ) - val_loss = validate(model, val_loader, device, cfg.data.image_size) + val_loss = validate(model, val_loader, device, + cfg.data.image_size, cfg.data.frame_len) scheduler.step() print(f" train_loss={train_loss:.6f} val_loss={val_loss:.6f}") diff --git a/actpred/test_actpred.py b/actpred/test_actpred.py index 370eb65..a92a6ad 100644 --- a/actpred/test_actpred.py +++ b/actpred/test_actpred.py @@ -79,45 +79,57 @@ class TestParseActionHistory(unittest.TestCase): class TestActionPredictor(unittest.TestCase): - """测试模型前向传播""" + """测试多帧注意力模型前向传播""" - def test_forward_simple_mlp(self): - data_cfg = DataConfig(image_size=64, action_dim=18, history_len=5) - model_cfg = ModelConfig(cnn_backbone="simple", history_encoder="mlp") + def test_forward_multi_frame(self): + data_cfg = DataConfig(image_size=64, action_dim=18, history_len=5, frame_len=3) + model_cfg = ModelConfig(cnn_backbone="simple") model = ActionPredictor(model_cfg, data_cfg) model.eval() - image = torch.randn(2, 3, 64, 64) + images = torch.randn(2, 3, 3, 64, 64) history = torch.randn(2, 5, 18) with torch.no_grad(): - out = model(image, history) + out = model(images, history) self.assertEqual(out.shape, (2, 18)) - def test_forward_simple_gru(self): - data_cfg = DataConfig(image_size=64, action_dim=18, history_len=5) - model_cfg = ModelConfig(cnn_backbone="simple", history_encoder="gru") + def test_forward_single_frame_degenerate(self): + """frame_len=1 退化为单帧输入也应正常前向""" + data_cfg = DataConfig(image_size=32, action_dim=4, history_len=3, frame_len=1) + model_cfg = ModelConfig(cnn_backbone="simple") model = ActionPredictor(model_cfg, data_cfg) model.eval() - image = torch.randn(2, 3, 64, 64) - history = torch.randn(2, 5, 18) + images = torch.randn(2, 1, 3, 32, 32) + history = torch.randn(2, 3, 4) with torch.no_grad(): - out = model(image, history) - self.assertEqual(out.shape, (2, 18)) + out = model(images, history) + self.assertEqual(out.shape, (2, 4)) def test_forward_dynamic_batch(self): - data_cfg = DataConfig(image_size=32, action_dim=4, history_len=3) - model_cfg = ModelConfig(cnn_backbone="simple", history_encoder="mlp") + data_cfg = DataConfig(image_size=32, action_dim=4, history_len=3, frame_len=2) + model_cfg = ModelConfig(cnn_backbone="simple") model = ActionPredictor(model_cfg, data_cfg) model.eval() for batch_size in [1, 4, 8]: - image = torch.randn(batch_size, 3, 32, 32) + images = torch.randn(batch_size, 2, 3, 32, 32) history = torch.randn(batch_size, 3, 4) with torch.no_grad(): - out = model(image, history) + out = model(images, history) self.assertEqual(out.shape, (batch_size, 4)) + def test_custom_attn_config(self): + """自定义 attn_dim/heads/layers 应可构建并前向""" + data_cfg = DataConfig(image_size=32, action_dim=4, history_len=3, frame_len=2) + model_cfg = ModelConfig(attn_dim=64, attn_heads=2, attn_layers=1) + model = ActionPredictor(model_cfg, data_cfg) + model.eval() + + with torch.no_grad(): + out = model(torch.randn(1, 2, 3, 32, 32), torch.randn(1, 3, 4)) + self.assertEqual(out.shape, (1, 4)) + class TestDataRecorder(unittest.TestCase): """测试数据录制器""" @@ -169,17 +181,30 @@ class TestActionDataset(unittest.TestCase): def test_dataset_loading(self): with tempfile.TemporaryDirectory() as tmpdir: self._create_dummy_episode(Path(tmpdir), n_frames=10, action_dim=4) - ds = ActionDataset(tmpdir, history_len=3, action_dim=4, image_size=32) - self.assertEqual(len(ds), 7) # 10 - 3 = 7 samples + ds = ActionDataset(tmpdir, history_len=3, action_dim=4, image_size=32, + frame_len=3) + # 采样起点 max(history_len, frame_len-1) = 3 → 10 - 3 = 7 samples + self.assertEqual(len(ds), 7) sample = ds[0] - self.assertIn("image", sample) + self.assertIn("images", sample) self.assertIn("action_history", sample) self.assertIn("target", sample) - self.assertEqual(sample["image"].shape, (3, 32, 32)) + self.assertEqual(sample["images"].shape, (3, 3, 32, 32)) self.assertEqual(sample["action_history"].shape, (3, 4)) self.assertEqual(sample["target"].shape, (4,)) + def test_frame_len_gt_history_shifts_start(self): + """frame_len-1 > history_len 时采样起点由帧窗口决定""" + with tempfile.TemporaryDirectory() as tmpdir: + self._create_dummy_episode(Path(tmpdir), n_frames=10, action_dim=4) + ds = ActionDataset(tmpdir, history_len=3, action_dim=4, image_size=32, + frame_len=6) + # 起点 max(3, 5) = 5 → 10 - 5 = 5 samples + self.assertEqual(len(ds), 5) + sample = ds[0] + self.assertEqual(sample["images"].shape, (6, 3, 32, 32)) + def test_episode_count(self): with tempfile.TemporaryDirectory() as tmpdir: self._create_dummy_episode(Path(tmpdir)) @@ -241,17 +266,24 @@ class TestConfig(unittest.TestCase): self.assertEqual(cfg.data.image_size, 64) self.assertEqual(cfg.data.action_dim, 18) self.assertEqual(cfg.data.history_len, 5) + self.assertEqual(cfg.data.frame_len, 3) self.assertEqual(cfg.model.cnn_backbone, "simple") - self.assertEqual(cfg.model.history_encoder, "mlp") + self.assertEqual(cfg.model.attn_dim, 128) + self.assertEqual(cfg.model.attn_heads, 4) + self.assertEqual(cfg.model.attn_layers, 2) self.assertEqual(cfg.train.epochs, 100) def test_apply_overrides(self): parser = build_argparser("test") - args = parser.parse_args(["--epochs", "50", "--lr", "0.01", "--backbone", "mobilenet"]) + args = parser.parse_args(["--epochs", "50", "--lr", "0.01", + "--backbone", "mobilenet", + "--frame-len", "4", "--attn-dim", "64"]) cfg = apply_overrides(Config(), args) self.assertEqual(cfg.train.epochs, 50) self.assertEqual(cfg.train.lr, 0.01) self.assertEqual(cfg.model.cnn_backbone, "mobilenet") + self.assertEqual(cfg.data.frame_len, 4) + self.assertEqual(cfg.model.attn_dim, 64) def test_normalize_actions_flag(self): """--normalize-actions 和 --no-normalize-actions 都应可用""" @@ -421,18 +453,23 @@ class TestValidateConfigValues(unittest.TestCase): def test_out_of_range_or_non_int_raises(self): for field, value in [("image_size", 99999), ("image_size", 0), - ("history_len", -1), ("action_dim", "18")]: + ("history_len", -1), ("frame_len", 0), + ("frame_len", 99999), ("action_dim", "18")]: cfg = Config() setattr(cfg.data, field, value) with self.assertRaises(ValueError, msg=f"{field}={value}"): validate_config_values(cfg.data, cfg.model) def test_model_field_out_of_range_raises(self): - """model 字段越界(fused_hidden_dim 超限、dropout 越界)应抛 ValueError""" + """model 字段越界(fused_hidden_dim 超限、dropout 越界、attn_* 越界)应抛 ValueError""" for field, value in [("fused_hidden_dim", 10 ** 9), ("fused_hidden_dim", 0), ("cnn_feature_dim", 99999), - ("action_hidden_dim", -1), + ("attn_dim", 0), + ("attn_dim", 99999), + ("attn_heads", 0), + ("attn_heads", 128), + ("attn_layers", 0), ("dropout", 1.5), ("dropout", -0.1)]: cfg = Config() @@ -440,6 +477,14 @@ class TestValidateConfigValues(unittest.TestCase): with self.assertRaises(ValueError, msg=f"{field}={value}"): validate_config_values(cfg.data, cfg.model) + def test_attn_dim_not_divisible_by_heads_raises(self): + """attn_dim 不能被 attn_heads 整除应抛 ValueError""" + cfg = Config() + cfg.model.attn_dim = 130 + cfg.model.attn_heads = 4 + with self.assertRaisesRegex(ValueError, "整除"): + validate_config_values(cfg.data, cfg.model) + def test_train_field_out_of_range_raises(self): """train 字段越界(lr 为负、val_split>1 等)应抛 ValueError""" for field, value in [("lr", -1.0), ("lr", 0.0), @@ -474,6 +519,15 @@ class TestParseCheckpointConfig(unittest.TestCase): with self.assertRaisesRegex(ValueError, "evil_extra_key"): parse_checkpoint_config(d) + def test_legacy_checkpoint_config_rejected(self): + """旧架构 checkpoint(含 history_encoder/action_hidden_dim)应拒载""" + cfg = Config() + d = cfg.to_dict() + d["model"]["history_encoder"] = "mlp" + d["model"]["action_hidden_dim"] = 64 + with self.assertRaisesRegex(ValueError, "history_encoder"): + parse_checkpoint_config(d) + def test_missing_section_raises(self): cfg = Config() d = cfg.to_dict() @@ -519,11 +573,13 @@ class TestActionPredictorInference(unittest.TestCase): """测试推理封装端到端:checkpoint 保存 → 加载 → 推理""" @staticmethod - def _make_checkpoint(tmpdir, normalize_actions=False, image_size=32): + def _make_checkpoint(tmpdir, normalize_actions=False, image_size=32, + frame_len=3): cfg = Config() cfg.data.image_size = image_size cfg.data.action_dim = 4 cfg.data.history_len = 3 + cfg.data.frame_len = frame_len cfg.data.normalize_actions = normalize_actions model = ActionPredictor(cfg.model, cfg.data) model.eval() @@ -533,17 +589,19 @@ class TestActionPredictorInference(unittest.TestCase): return path, model, cfg def test_end_to_end_predict(self): - """加载 checkpoint 后推理输出形状正确、数值有限且可复现""" + """加载 checkpoint 后多帧推理输出形状正确、数值有限且可复现""" with tempfile.TemporaryDirectory() as tmpdir: path, _, _ = self._make_checkpoint(tmpdir) predictor = ActionPredictorInference(path, device="cpu") self.assertEqual(predictor.action_dim, 4) self.assertEqual(predictor.history_len, 3) + self.assertEqual(predictor.frame_len, 3) - image = np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8) + images = [np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8) + for _ in range(3)] history = np.random.randn(3, 4).astype(np.float32) - pred1 = predictor.predict(image, history) - pred2 = predictor.predict(image, history) + pred1 = predictor.predict(images, history) + pred2 = predictor.predict(images, history) self.assertEqual(pred1.shape, (4,)) self.assertTrue(np.all(np.isfinite(pred1))) np.testing.assert_array_equal(pred1, pred2) # eval 模式确定性 @@ -563,9 +621,18 @@ class TestActionPredictorInference(unittest.TestCase): with tempfile.TemporaryDirectory() as tmpdir: path, _, _ = self._make_checkpoint(tmpdir) predictor = ActionPredictorInference(path, device="cpu") - image = np.zeros((32, 32, 3), dtype=np.uint8) + images = [np.zeros((32, 32, 3), dtype=np.uint8)] * 3 with self.assertRaisesRegex(ValueError, "action_history"): - predictor.predict(image, np.zeros((2, 4), dtype=np.float32)) + predictor.predict(images, np.zeros((2, 4), dtype=np.float32)) + + def test_wrong_frame_count_raises(self): + """传入帧数与 frame_len 不一致应抛 ValueError""" + with tempfile.TemporaryDirectory() as tmpdir: + path, _, _ = self._make_checkpoint(tmpdir) + predictor = ActionPredictorInference(path, device="cpu") + images = [np.zeros((32, 32, 3), dtype=np.uint8)] * 2 + with self.assertRaisesRegex(ValueError, "帧"): + predictor.predict(images, np.zeros((3, 4), dtype=np.float32)) def test_missing_action_stats_raises(self): """启用归一化但缺 action_stats.npy 应抛 FileNotFoundError""" @@ -594,9 +661,9 @@ class TestActionPredictorInference(unittest.TestCase): image = np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8) history = np.random.randn(3, 4).astype(np.float32) pred_raw = ActionPredictorInference( - raw_path, device="cpu").predict(image, history) + raw_path, device="cpu").predict([image] * 3, history) pred_norm = ActionPredictorInference( - norm_path, device="cpu").predict(image, history) + norm_path, device="cpu").predict([image] * 3, history) np.testing.assert_allclose(pred_norm, pred_raw, rtol=1e-5, atol=1e-6) def test_stats_shape_mismatch_raises(self): diff --git a/actpred/train.py b/actpred/train.py index b8d61aa..50b8ffe 100644 --- a/actpred/train.py +++ b/actpred/train.py @@ -25,18 +25,18 @@ def train_one_epoch(model, loader, optimizer, device, max_grad_norm: float = Non total_loss = 0.0 count = 0 for batch in tqdm(loader, desc=" train", leave=False): - image = batch["image"].to(device) + images = batch["images"].to(device) history = batch["action_history"].to(device) target = batch["target"].to(device) optimizer.zero_grad() - pred = model(image, history) + pred = model(images, 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) + total_loss += loss.item() * images.size(0) + count += images.size(0) return total_loss / max(count, 1) @@ -46,13 +46,13 @@ def validate(model, loader, device): total_loss = 0.0 count = 0 for batch in tqdm(loader, desc=" val ", leave=False): - image = batch["image"].to(device) + images = batch["images"].to(device) history = batch["action_history"].to(device) target = batch["target"].to(device) - pred = model(image, history) + pred = model(images, history) loss = torch.nn.functional.mse_loss(pred, target) - total_loss += loss.item() * image.size(0) - count += image.size(0) + total_loss += loss.item() * images.size(0) + count += images.size(0) if count == 0: return float("inf") # 空验证集不参与比较,避免假 best return total_loss / count @@ -83,7 +83,7 @@ def main(): train_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, train_indices) + cfg.data.action_stats_path, train_indices, cfg.data.frame_len) 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") @@ -93,7 +93,7 @@ def main(): 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) + cfg.data.action_stats_path, val_indices, cfg.data.frame_len) print(f"Samples: {len(train_ds)} train, {len(val_ds)} val") if len(train_ds) == 0: