示例脚本跑不通、输出帧和输入画面对不上，多数时候不是模型本身的问题，而是喂进去的张量形状和仓库约定不一致：维度顺序反了、采样帧数差一帧、通道维被当成时间维。比较省事的顺序是——先把输入输出约定固定成一份可执行的契约，在调用前向之前用断言把形状错误挡住，再去调参或看画质。

Cosmos 3 Edge 这类「输入图像与动作序列、输出预测帧」的链路，第一步应当是确认张量维度语义，而不是先跑一次看结果。建议动作：读仓库代码确认排列，按约定组装 batch，逐层打印输出形状，把 assert 前置，最后把约定写进适配层并配单元测试。形状取值以仓库配置和实际打印为准，不同版本脚本可能不一致，需要结合自己的环境确认。

从示例脚本里读出输入张量的维度顺序

维度顺序不要靠记忆推。找到仓库里真正被调用的推理脚本或模型的 forward 定义，在组装完输入、发起模型调用之前插一行打印，把维度个数和每一维大小一起打出来——只打 shape 不打维度名，很容易把 B 和 T 看反。

video, actions = build_inputs(...)   # 或从 dataloader 取一个 batch
print("video:", tuple(video.shape), video.dtype)
print("actions:", tuple(actions.shape), actions.dtype)
out = model(video, actions)

视频张量常见的两种排列是：(B, T, C, H, W) 与 (B, C, T, H, W)。前者把帧维放在通道维前面，采样和切片时 video[:, t] 直接拿到一帧图像；后者把时间维塞进通道位置，卷积层容易直接吃下去，但你在取单帧时就得写成 video[:, :, t]。解码器往往先给到 (T, H, W, C)，需要补 batch 维并 permute，这一步最容易漏。

• 看 forward 签名和内部第一次 reshape / permute 的位置，那里通常写着作者的假设。
• 看配置文件里的 num_frames、in_channels、height、width，和打印结果逐项对照。
• 看 dataloader 的 collate，确认 batch 维是 collate 补的还是脚本自己 unsqueeze 补的。

上述只是常见情况，最终以仓库代码为准：同一份示例在不同分支上变更过排列的情况并不少见，打印一次比读三遍注释可靠。

把一段视频和动作序列组装成 batch

下面是一段可替换的组装骨架，输入是解码得到的帧列表和动作数组，输出是符合 (B, T, C, H, W) 约定的张量。每一步都带形状打印，方便定位在哪一步出错。

import torch

def build_batch(frames, actions, num_frames=9, height=256, width=256, device="cpu"):
    # frames: 解码结果列表，元素为 HWC 排列的图像数组
    idx = torch.linspace(0, len(frames) - 1, num_frames).long()
    sampled = [frames[i] for i in idx]
    video = torch.stack([torch.as_tensor(f) for f in sampled])   # (T, H, W, C)
    print("sampled:", tuple(video.shape))

    video = video.permute(0, 3, 1, 2).float() / 255.0           # (T, C, H, W)
    print("permute:", tuple(video.shape))

    video = video.unsqueeze(0)                                  # (B=1, T, C, H, W)
    print("batch video:", tuple(video.shape))

    act = torch.as_tensor(actions, dtype=torch.float32)
    print("raw actions:", tuple(act.shape))                     # 期望 (T_act, A)
    if act.dim() == 2:
        act = act[:num_frames]                                  # 与帧数对齐
    act = act.unsqueeze(0)                                      # (B, T_act, A)
    print("batch actions:", tuple(act.shape))
    return video.to(device), act.to(device)

帧采样要保持时序，等间隔取比随机取更容易复现；动作序列与帧的对应关系可能有步长（一个动作对应多帧或反之），这个比例需要从仓库的动作处理代码里确认，不要默认一一对应。若动作长度不足，宁可截断到与帧对齐再喂，也不要在模型内部被广播成奇怪形状。

跑前向并打印每层输出形状

形状对不上时，问题常常出在中间某一层把时间维和通道维合并了。用 forward hook 逐层打印比在源码里到处插 print 更安全，也方便一次跑完看全貌。

def shape_hook(name):
    def fn(mod, inp, out):
        o = out[0] if isinstance(out, (tuple, list)) else out
        print(f"{name:32s} {type(mod).__name__:16s} out={tuple(o.shape)}")
    return fn

handles = [m.register_forward_hook(shape_hook(n))
           for n, m in model.named_modules()
           if len(list(m.children())) == 0]        # 只挂叶子模块，避免刷屏
try:
    out = model(video, actions)
finally:
    for h in handles:
        h.remove()
print("final:", tuple(out.shape))

把首次运行的结果整理成一张逐层记录表，后续改动可以直接对照，不必每次重跑再肉眼扫日志。

从第一个不匹配的层往回看它的输入，通常就能确定是上游 permute、reshape 还是 concat 把维度顺序改掉了。只记录「哪一层开始不一样」比记录全部层更省时间。

用形状断言代替肉眼比对

把维度检查写成 assert，放在构造输入之后、调用模型之前。这样错误在第一时间抛出，报错信息里直接带着实际形状和期望值，不用再翻日志猜。

def check_inputs(video, actions, cfg):
    assert video.dim() == 5, \
        f"video 需要 5 维 (B,T,C,H,W)，实际 dim={video.dim()} shape={tuple(video.shape)}"
    B, T, C, H, W = video.shape
    assert T == cfg.num_frames, \
        f"帧数不匹配: T={T}, cfg.num_frames={cfg.num_frames}"
    assert C == cfg.in_channels, \
        f"通道数不匹配: C={C}, 期望 {cfg.in_channels}"
    assert (H, W) == (cfg.height, cfg.width), \
        f"分辨率不匹配: 实际 {(H, W)}，期望 {(cfg.height, cfg.width)}"
    assert actions.dim() == 3 and actions.shape[0] == B, \
        f"actions 需要 (B,T_act,A)，实际 {tuple(actions.shape)}"

看报错时优先看三个字段：实际 shape 元组、期望值与实际值是否只差一维、以及 dim() 是不是差 1。如果实际形状比期望少一维，多半是漏了 batch 的 unsqueeze；如果某一维恰好等于帧数却出现在通道位置，基本可以判定是排列顺序问题而不是数据问题。

把形状约定写进项目的输入适配层

把上面几步收进一个小的适配函数，调用方只传原始帧和动作，形状对齐由适配层统一负责。这样后续换模型或换数据源，改动集中在一处。

class InputAdapter:
    def __init__(self, num_frames, height, width, in_channels=3):
        self.num_frames, self.height, self.width = num_frames, height, width
        self.in_channels = in_channels

    def __call__(self, frames, actions):
        assert len(frames) >= 1, "frames 不能为空"
        video, act = build_batch(frames, actions, self.num_frames,
                                 self.height, self.width)
        check_inputs(video, act, self)
        return {"video": video, "actions": act}

配套的单元测试不必依赖真实视频，用随机张量把契约钉住即可。

import pytest, torch
from adapter import InputAdapter

def test_shape_contract():
    a = InputAdapter(num_frames=4, height=64, width=64)
    frames = [torch.randint(0, 255, (64, 64, 3), dtype=torch.uint8) for _ in range(10)]
    out = a(frames, torch.zeros(4, 7))
    assert out["video"].shape == (1, 4, 3, 64, 64)
    assert out["actions"].shape[0] == 1

def test_bad_action_dim_raises():
    a = InputAdapter(num_frames=4, height=64, width=64)
    frames = [torch.zeros(64, 64, 3, dtype=torch.uint8) for _ in range(4)]
    with pytest.raises(AssertionError):
        a(frames, torch.zeros(4))          # 少一维，应被断言拦下

def test_frame_shortage_raises():
    a = InputAdapter(num_frames=8, height=64, width=64)
    frames = [torch.zeros(64, 64, 3, dtype=torch.uint8) for _ in range(2)]
    with pytest.raises(Exception):
        a(frames, torch.zeros(8, 7))

测试里只断言形状与异常类型，不断言具体像素或性能，避免环境差异导致误报。适配层里的帧数、分辨率、通道数建议从配置读取而不是写死，换分辨率或换帧数时改配置并重跑这几个测试，就能确认契约没有被破坏。