From 91b30040c7338235e2c2067da1faa5666520f024 Mon Sep 17 00:00:00 2001 From: qing <1776177662@qq.com> Date: Thu, 13 Aug 2026 02:19:30 +0800 Subject: [PATCH] =?UTF-8?q?refactor(actpred):=20=E6=A8=A1=E5=9E=8B?= =?UTF-8?q?=E8=BE=93=E5=85=A5=E5=B0=BA=E5=AF=B8=20image=5Fsize=20=E2=86=92?= =?UTF-8?q?=20image=5Fwidth/image=5Fheight=EF=BC=8C=E9=BB=98=E8=AE=A4=2019?= =?UTF-8?q?20=C3=971080?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit config/model/dataset/train/pretrain/predict/export 全链路适配非方形输入， hailo_infer/onnx_infer 从模型元数据动态推导宽高；旧 224 checkpoint 配置键 不兼容（重训后替换）。Pi 5 CPU 冒烟验证 1080p 训练可行（~4.75s/步 batch=2）。 --- actpred/README.md | 5 +++-- actpred/config.py | 12 ++++++++---- actpred/dataset.py | 8 +++++--- actpred/export_onnx.py | 2 +- actpred/hailo_infer.py | 24 ++++++++++++------------ actpred/model.py | 7 ++++--- actpred/onnx_infer.py | 12 ++++++------ actpred/predict.py | 3 ++- actpred/pretrain_mt.py | 14 ++++++++------ actpred/test_actpred.py | 38 ++++++++++++++++++++++---------------- actpred/train.py | 4 ++-- 11 files changed, 73 insertions(+), 56 deletions(-) diff --git a/actpred/README.md b/actpred/README.md index 2461fb7..441aa32 100644 --- a/actpred/README.md +++ b/actpred/README.md @@ -9,7 +9,7 @@ Transformer（pre-LN）自注意力融合，取 [PRED] token 回归下一步动 ## 架构 ``` -F 帧图像 (F×3×224×224) ─► 共享 CNN 编码器 ─► F 个 token ─► proj → d_model ┐ +F 帧图像 (F×3×H×W) ─► 共享 CNN 编码器 ─► F 个 token ─► proj → d_model ┐ ├─► [img×F, act×H, PRED] + 位置编码 动作历史 (H×D) ─► Linear proj → d_model ─► H 个 token ────────┤ │ 可学习 [PRED] token ────────────────────────────────────────┘ ▼ @@ -110,7 +110,8 @@ recorder.flush() | 参数 | 默认值 | 说明 | |------|--------|------| -| `--image-size` | 224 | 图像缩放尺寸 | +| `--image-width` | 1920 | 图像缩放宽度 | +| `--image-height` | 1080 | 图像缩放高度 | | `--action-dim` | 18 | 动作向量维度（对齐 bot4cj 18 列：10 舵机 + 2 云台 + 扬声器频率 + 延时 + 4 预留）| | `--history-len` | 5 | 历史动作窗口长度 | | `--frame-len` | 3 | 输入图像帧数（含当前帧，旧帧在前）| diff --git a/actpred/config.py b/actpred/config.py index 85c1e98..50aff32 100644 --- a/actpred/config.py +++ b/actpred/config.py @@ -5,7 +5,8 @@ from typing import Optional @dataclass class DataConfig: - image_size: int = 224 + image_width: int = 1920 + image_height: int = 1080 action_dim: int = 18 history_len: int = 5 frame_len: int = 3 @@ -59,7 +60,8 @@ class Config: def build_argparser(description: str = "") -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description=description) - parser.add_argument("--image-size", type=int, default=None) + parser.add_argument("--image-width", type=int, default=None) + parser.add_argument("--image-height", 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) @@ -99,7 +101,8 @@ def validate_config_values(data_cfg: DataConfig, model_cfg: ModelConfig, 也防止训练 CLI 越界参数（如 --history-len 0 / --lr -1）静默进入训练。 """ checks = [ - ("image_size", data_cfg.image_size, 1, 1024, int), + ("image_width", data_cfg.image_width, 1, 4096, int), + ("image_height", data_cfg.image_height, 1, 4096, 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), @@ -173,7 +176,8 @@ def parse_checkpoint_config(cfg_dict: dict) -> Config: def apply_overrides(cfg: Config, args: argparse.Namespace) -> Config: mapping = { - "image_size": ("data", "image_size"), + "image_width": ("data", "image_width"), + "image_height": ("data", "image_height"), "action_dim": ("data", "action_dim"), "history_len": ("data", "history_len"), "frame_len": ("data", "frame_len"), diff --git a/actpred/dataset.py b/actpred/dataset.py index 489198d..db3f25b 100644 --- a/actpred/dataset.py +++ b/actpred/dataset.py @@ -16,16 +16,18 @@ def _natural_key(p): class ActionDataset(Dataset): def __init__(self, data_dir: str, history_len: int, action_dim: int, - image_size: int = 224, normalize_actions: bool = False, + image_width: int = 1920, image_height: int = 1080, + normalize_actions: bool = False, 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.image_width = image_width + self.image_height = image_height self.frame_len = frame_len self.transform = transforms.Compose([ - transforms.Resize((image_size, image_size)), + transforms.Resize((image_height, image_width)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) diff --git a/actpred/export_onnx.py b/actpred/export_onnx.py index 55ab562..9c36c00 100644 --- a/actpred/export_onnx.py +++ b/actpred/export_onnx.py @@ -21,7 +21,7 @@ def main(): model.eval() dummy_images = torch.randn( - 1, data_cfg.frame_len, 3, data_cfg.image_size, data_cfg.image_size) + 1, data_cfg.frame_len, 3, data_cfg.image_height, data_cfg.image_width) dummy_history = torch.randn(1, data_cfg.history_len, data_cfg.action_dim) batch = Dim("batch") diff --git a/actpred/hailo_infer.py b/actpred/hailo_infer.py index 8254d58..1f0155e 100644 --- a/actpred/hailo_infer.py +++ b/actpred/hailo_infer.py @@ -5,8 +5,8 @@ HEF 由 actpred ONNX 经 Hailo DFC 编译生成（量化校准后）。本封装不硬编码 vstream 名称/形状/量化参数，全部从 HEF 的 input/output vstream info 动态推导： - 图像输入: (1, F, 3, S, S) → NHWC uint8 (1, F, S, S, 3) - 或 (1, S, S, F*3) → 帧沿通道拼接 + 图像输入: (1, F, 3, H, W) → NHWC uint8 (1, F, H, W, 3) + 或 (1, H, W, F*3) → 帧沿通道拼接 动作输入: (1, H, D) float32，uint8 边缘时按 quant_info 量化 动作输出: (1, D)，uint8 边缘时按 quant_info 反量化 @@ -111,15 +111,15 @@ class HailoPredictor: def _derive_dims(self): shape = tuple(self._image_info.shape) if len(shape) == 5: - _, self.frame_len, c, self.image_size, _ = shape + _, self.frame_len, c, self.image_height, self.image_width = shape if c != 3: raise ValueError(f"图像输入通道维应为 3，实际形状 {shape}") self._image_layout = "ncfhw" - else: # (1, S, S, F*3) NHWC 帧拼通道 - _, self.image_size, s2, c = shape - if s2 != self.image_size or c % 3 != 0 or c == 3: + else: # (1, H, W, F*3) NHWC 帧拼通道 + _, self.image_height, self.image_width, c = shape + if c % 3 != 0 or c == 3: raise ValueError( - f"无法从 4 维图像输入推导帧数（期望 (1,S,S,F*3)），" + f"无法从 4 维图像输入推导帧数（期望 (1,H,W,F*3)），" f"实际形状 {shape}") self.frame_len = c // 3 self._image_layout = "nhwc_fused" @@ -206,15 +206,15 @@ class HailoPredictor: def _preprocess(self, image) -> np.ndarray: img = Image.fromarray(image) if isinstance(image, np.ndarray) else image img = img.convert("RGB").resize( - (self.image_size, self.image_size), Image.BILINEAR) + (self.image_width, self.image_height), Image.BILINEAR) return np.asarray(img, dtype=np.uint8) # HWC uint8 def _image_buffer(self, frames: np.ndarray) -> np.ndarray: - """(F, S, S, 3) → HEF 图像输入缓冲区（按编译后布局）""" + """(F, H, W, 3) → HEF 图像输入缓冲区（按编译后布局）""" if self._image_layout == "ncfhw": - return frames.transpose(0, 3, 1, 2)[np.newaxis] # (1,F,3,S,S) - # nhwc_fused: 帧沿通道拼接 (1, S, S, F*3) - return frames.reshape(1, self.image_size, self.image_size, + return frames.transpose(0, 3, 1, 2)[np.newaxis] # (1,F,3,H,W) + # nhwc_fused: 帧沿通道拼接 (1, H, W, F*3) + return frames.reshape(1, self.image_height, self.image_width, self.frame_len * 3) diff --git a/actpred/model.py b/actpred/model.py index 9eaa817..50b229c 100644 --- a/actpred/model.py +++ b/actpred/model.py @@ -127,14 +127,15 @@ class ActionPredictor(nn.Module): 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._image_width = data_cfg.image_width + self._image_height = data_cfg.image_height self._cnn_feature_dim = model_cfg.cnn_feature_dim def forward(self, images: torch.Tensor, action_history: torch.Tensor) -> torch.Tensor: - # images: (B, F, 3, S, S) action_history: (B, H, D) + # images: (B, F, 3, H, W) action_history: (B, H_len, D) b = images.size(0) frames = images.reshape( - b * self._frame_len, 3, self._image_size, self._image_size) + b * self._frame_len, 3, self._image_height, self._image_width) img_feat = self.image_encoder(frames).reshape( b, self._frame_len, self._cnn_feature_dim) img_tokens = self.img_proj(img_feat) diff --git a/actpred/onnx_infer.py b/actpred/onnx_infer.py index e53be78..4f56f4b 100644 --- a/actpred/onnx_infer.py +++ b/actpred/onnx_infer.py @@ -2,7 +2,7 @@ 接口对齐 predict.ActionPredictorInference：predict(images, action_history)。 -模型维度从 ONNX session 输入元数据动态推导（images: [1,F,3,S,S]， +模型维度从 ONNX session 输入元数据动态推导（images: [1,F,3,H,W]， action_history: [1,H,D]），无需 checkpoint。 图像预处理复现训练时的 dataset.py 流水线：Resize → /255 → ImageNet 归一化。 @@ -55,7 +55,7 @@ class OnnxPredictor: inputs = {i.name: [d if isinstance(d, int) else -1 for d in i.shape] for i in self._session.get_inputs()} img_shape, hist_shape = self._locate_inputs(inputs) - _, self.frame_len, _, self.image_size, _ = img_shape + _, self.frame_len, _, self.image_height, self.image_width = img_shape _, self.history_len, self.action_dim = hist_shape self._image_name = next(n for n, s in inputs.items() if len(s) == 5) self._history_name = next(n for n, s in inputs.items() if len(s) == 3) @@ -82,9 +82,9 @@ class OnnxPredictor: raise ValueError( "ONNX 输入应为 1 个 5 维图像 + 1 个 3 维动作历史，" f"实际: {inputs}") - if img[0][2] != 3 or img[0][3] != img[0][4]: + if img[0][2] != 3: raise ValueError( - f"图像输入形状应为 [1,F,3,S,S]，实际为 {img[0]}") + f"图像输入形状应为 [1,F,3,H,W]，实际为 {img[0]}") return img[0], hist[0] # ---------------- 生命周期 ---------------- @@ -134,10 +134,10 @@ class OnnxPredictor: return result def _preprocess(self, image) -> np.ndarray: - """Resize → /255 → ImageNet 归一化 → (3, S, S) float32""" + """Resize → /255 → ImageNet 归一化 → (3, H, W) float32""" img = Image.fromarray(image) if isinstance(image, np.ndarray) else image img = img.convert("RGB").resize( - (self.image_size, self.image_size), Image.BILINEAR) + (self.image_width, self.image_height), Image.BILINEAR) x = np.asarray(img, dtype=np.float32) / 255.0 # HWC x = x.transpose(2, 0, 1) # CHW return (x - IMAGENET_MEAN[:, None, None]) / IMAGENET_STD[:, None, None] diff --git a/actpred/predict.py b/actpred/predict.py index f09cf3c..f547826 100644 --- a/actpred/predict.py +++ b/actpred/predict.py @@ -31,7 +31,8 @@ class ActionPredictorInference: self.frame_len = self._data_cfg.frame_len self._transform = transforms.Compose([ - transforms.Resize((self._data_cfg.image_size, self._data_cfg.image_size)), + transforms.Resize((self._data_cfg.image_height, + self._data_cfg.image_width)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) diff --git a/actpred/pretrain_mt.py b/actpred/pretrain_mt.py index e3b06b4..1948392 100644 --- a/actpred/pretrain_mt.py +++ b/actpred/pretrain_mt.py @@ -76,7 +76,8 @@ class MTActionDataset(Dataset): } -def pretrain_one_epoch(model, loader, optimizer, device, image_size, frame_len, +def pretrain_one_epoch(model, loader, optimizer, device, image_width, + image_height, frame_len, max_grad_norm=None, noise_std=0.0): model.train() total_loss = 0.0 @@ -90,7 +91,7 @@ def pretrain_one_epoch(model, loader, optimizer, device, image_size, frame_len, target = target + torch.randn_like(target) * noise_std 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, + history.size(0), frame_len, 3, image_height, image_width, device=device, ) optimizer.zero_grad() @@ -106,7 +107,7 @@ def pretrain_one_epoch(model, loader, optimizer, device, image_size, frame_len, @torch.no_grad() -def validate(model, loader, device, image_size, frame_len): +def validate(model, loader, device, image_width, image_height, frame_len): model.eval() total_loss = 0.0 count = 0 @@ -116,7 +117,7 @@ def validate(model, loader, device, image_size, frame_len): target = batch["target"].to(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, + history.size(0), frame_len, 3, image_height, image_width, device=device, ) pred = model(zero_images, history) @@ -239,11 +240,12 @@ 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.data.frame_len, + cfg.data.image_width, cfg.data.image_height, cfg.data.frame_len, cfg.train.max_grad_norm, args.noise_std, ) val_loss = validate(model, val_loader, device, - cfg.data.image_size, cfg.data.frame_len) + cfg.data.image_width, cfg.data.image_height, + 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 ca0a86d..d2f2bd1 100644 --- a/actpred/test_actpred.py +++ b/actpred/test_actpred.py @@ -84,7 +84,7 @@ class TestActionPredictor(unittest.TestCase): """测试多帧注意力模型前向传播""" def test_forward_multi_frame(self): - data_cfg = DataConfig(image_size=64, action_dim=18, history_len=5, frame_len=3) + data_cfg = DataConfig(image_width=64, image_height=64, action_dim=18, history_len=5, frame_len=3) model_cfg = ModelConfig(cnn_backbone="simple") model = ActionPredictor(model_cfg, data_cfg) model.eval() @@ -97,7 +97,7 @@ class TestActionPredictor(unittest.TestCase): 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) + data_cfg = DataConfig(image_width=32, image_height=32, action_dim=4, history_len=3, frame_len=1) model_cfg = ModelConfig(cnn_backbone="simple") model = ActionPredictor(model_cfg, data_cfg) model.eval() @@ -109,7 +109,7 @@ class TestActionPredictor(unittest.TestCase): self.assertEqual(out.shape, (2, 4)) def test_forward_dynamic_batch(self): - data_cfg = DataConfig(image_size=32, action_dim=4, history_len=3, frame_len=2) + data_cfg = DataConfig(image_width=32, image_height=32, action_dim=4, history_len=3, frame_len=2) model_cfg = ModelConfig(cnn_backbone="simple") model = ActionPredictor(model_cfg, data_cfg) model.eval() @@ -123,7 +123,7 @@ class TestActionPredictor(unittest.TestCase): 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) + data_cfg = DataConfig(image_width=32, image_height=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() @@ -183,7 +183,8 @@ 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, + ds = ActionDataset(tmpdir, history_len=3, action_dim=4, + image_width=32, image_height=32, frame_len=3) # 采样起点 max(history_len, frame_len-1) = 3 → 10 - 3 = 7 samples self.assertEqual(len(ds), 7) @@ -200,7 +201,8 @@ class TestActionDataset(unittest.TestCase): """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, + ds = ActionDataset(tmpdir, history_len=3, action_dim=4, + image_width=32, image_height=32, frame_len=6) # 起点 max(3, 5) = 5 → 10 - 5 = 5 samples self.assertEqual(len(ds), 5) @@ -265,7 +267,8 @@ class TestConfig(unittest.TestCase): def test_default_config(self): cfg = Config() - self.assertEqual(cfg.data.image_size, 224) + self.assertEqual(cfg.data.image_width, 1920) + self.assertEqual(cfg.data.image_height, 1080) self.assertEqual(cfg.data.action_dim, 18) self.assertEqual(cfg.data.history_len, 5) self.assertEqual(cfg.data.frame_len, 3) @@ -480,7 +483,8 @@ class TestValidateConfigValues(unittest.TestCase): validate_config_values(cfg.data, cfg.model) # 不应抛异常 def test_out_of_range_or_non_int_raises(self): - for field, value in [("image_size", 99999), ("image_size", 0), + for field, value in [("image_width", 99999), ("image_width", 0), + ("image_height", 0), ("history_len", -1), ("frame_len", 0), ("frame_len", 99999), ("action_dim", "18")]: cfg = Config() @@ -601,10 +605,11 @@ class TestActionPredictorInference(unittest.TestCase): """测试推理封装端到端：checkpoint 保存 → 加载 → 推理""" @staticmethod - def _make_checkpoint(tmpdir, normalize_actions=False, image_size=32, - frame_len=3): + def _make_checkpoint(tmpdir, normalize_actions=False, image_width=32, + image_height=32, frame_len=3): cfg = Config() - cfg.data.image_size = image_size + cfg.data.image_width = image_width + cfg.data.image_height = image_height cfg.data.action_dim = 4 cfg.data.history_len = 3 cfg.data.frame_len = frame_len @@ -707,11 +712,11 @@ class TestActionPredictorInference(unittest.TestCase): with tempfile.TemporaryDirectory() as tmpdir: path, model, cfg = self._make_checkpoint(tmpdir) cfg_dict = cfg.to_dict() - cfg_dict["data"]["image_size"] = 99999 + cfg_dict["data"]["image_width"] = 99999 bad_path = os.path.join(tmpdir, "bad.pt") torch.save({"model_state_dict": model.state_dict(), "config": cfg_dict}, bad_path) - with self.assertRaisesRegex(ValueError, "image_size"): + with self.assertRaisesRegex(ValueError, "image_width"): ActionPredictorInference(bad_path, device="cpu") def test_unknown_config_key_in_checkpoint_raises(self): @@ -820,7 +825,8 @@ class TestOnnxPredictor(unittest.TestCase): @staticmethod def _export_small_onnx(tmpdir): cfg = Config() - cfg.data.image_size = 32 + cfg.data.image_width = 32 + cfg.data.image_height = 32 cfg.data.action_dim = 4 cfg.data.history_len = 3 cfg.data.frame_len = 3 @@ -847,7 +853,7 @@ class TestOnnxPredictor(unittest.TestCase): with OnnxPredictor(path) as p: self.assertEqual((p.frame_len, p.history_len, p.action_dim), (3, 3, 4)) - self.assertEqual(p.image_size, 32) + self.assertEqual((p.image_width, p.image_height), (32, 32)) def test_matches_torch_forward(self): from onnx_infer import OnnxPredictor @@ -1043,7 +1049,7 @@ class TestHailoPredictor(unittest.TestCase): history = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32) with HailoPredictor("fake.hef") as p: self.assertEqual((p.frame_len, p.history_len, p.action_dim), (2, 2, 3)) - self.assertEqual(p.image_size, 4) + self.assertEqual((p.image_width, p.image_height), (4, 4)) result = p.predict(self._frames(2), history) np.testing.assert_allclose(result, [1, 2, 3]) self.assertEqual(seen_shapes["image"], (1, 2, 3, 4, 4)) diff --git a/actpred/train.py b/actpred/train.py index 90ced88..b99956a 100644 --- a/actpred/train.py +++ b/actpred/train.py @@ -82,7 +82,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.image_width, cfg.data.image_height, cfg.data.normalize_actions, 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 不可信 @@ -92,7 +92,7 @@ def main(): 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.image_width, cfg.data.image_height, cfg.data.normalize_actions, cfg.data.action_stats_path, val_indices, cfg.data.frame_len) print(f"Samples: {len(train_ds)} train, {len(val_ds)} val")