这篇教程能做出什么
跟着走完，你会得到一个可运行的最小流程：拿一批无标注视频做自监督预训练，得到一个通用的视觉编码器；把这个编码器搬到 ARC 的彩色网格上，用每个任务自带的少量示例做测试时适配，最后跑一个评测脚本，得到逐任务的准确率。
"看猫片学 ARC"这个思路的直觉是：视频里天然包含物体恒存、遮挡、计数、对称、复制、移动这些结构。一只猫从椅子后面走过去再出现，模型要能"补全"中间那段；几个球滚来滚去，模型要学到"数量守恒"。这些先验和 ARC 要求的抽象规则高度重叠，而视频是免费的、海量的、不需要人工标注的。所以先看片、后解题，是一条值得试的迁移路线。
需要先说清楚预期：这条路线在公开评测上的绝对分数不会高，它的价值在于验证"视觉自监督预训练能带来可迁移的抽象先验"这个假设。教程里所有超参数都只是起点值，需要你按自己的数据和算力调整。
前置条件清单
- 一台带 GPU 的机器，能跑 224×224 分辨率的视频片段训练即可；纯 CPU 也能跑通流程，只是慢
- Python 环境，安装 PyTorch、torchvision、NumPy、Pillow、tqdm
- ffmpeg 命令行工具，用来抽帧
- 一批无标注视频：自己用手机拍的、公开动作识别数据集的视频都可以。数据集的使用条款和下载方式以各数据集官方页面为准
- ARC 的公开数据文件（JSON 格式，包含 train/test 两组示例对），获取方式以官方仓库当前说明为准
- 不需要预训练权重，一切从零开始；如果你想省时间，也可以先用公开的视觉预训练权重初始化，再跑本流程
第一步：搭好目录与配置
先固定目录结构，后面每一步的文件都有明确去处，排错时不会翻乱。
```text
arc-video-repro/
├── configs/
│ └── default.yaml
├── data/
│ ├── videos/ # 原始视频文件
│ ├── frames/ # 抽帧结果，每个视频一个子目录
│ └── arc/ # ARC 的 JSON 任务文件
├── checkpoints/
├── src/
│ ├── dataset.py # 视频片段数据集
│ ├── model.py # 自监督视频模型
│ ├── train_pretrain.py # 预训练入口
│ ├── arc_data.py # 网格与张量的互转
│ ├── transfer.py # 表征迁移与测试时适配
│ └── evaluate.py # 评测入口
└── requirements.txt
```
写下依赖清单，方便别人复现：
```text
torch
torchvision
numpy
pillow
tqdm
pyyaml
```
第二步：把无标注视频切成训练用的帧
自监督目标的输入是"视频片段"，也就是连续的若干帧。先把视频统一抽成帧序列存盘，训练时随取随用，比每次解码视频快很多。
```bash
每个视频抽成 4 fps 的帧，统一缩放到 224x224
mkdir -p data/frames
for f in data/videos/*.mp4; do
name=$(basename "$f" .mp4)
mkdir -p "data/frames/$name"
ffmpeg -i "$f" -vf "fps=4,scale=224:224" -q:v 3 "data/frames/$name/%06d.jpg"
done
```
帧率不用太高：ARC 里的规则大多是静态或定性变化，4 fps 已经能覆盖"物体移动、遮挡、复制"这类模式。分辨率保持能被 patch 尺寸整除（224 能被 16 整除），省去后面处理边界的麻烦。
写数据集类，每次返回一个定长片段：
```python
src/dataset.py
import os
import random
import torch
from torch.utils.data import Dataset
from torchvision import transforms
from PIL import Image
class ClipDataset(Dataset):
"""从每个帧目录里随机截取一段连续帧，返回 (T, 3, H, W) 张量。"""
def __init__(self, root, clip_len=8, size=224):
self.dirs = [
os.path.join(root, d)
for d in sorted(os.listdir(root))
if os.path.isdir(os.path.join(root, d))
]
self.clip_len = clip_len
self.tf = transforms.Compose([
transforms.Resize((size, size)),
transforms.ToTensor(), # 到 [0, 1]
])
def __len__(self):
return len(self.dirs)
def __getitem__(self, idx):
files = sorted(os.listdir(self.dirs[idx]))
k = self.clip_len
if len(files) < k:
files = files + [files[-1]] * (k - len(files))
start = random.randint(0, len(files) - k)
frames = [
self.tf(Image.open(os.path.join(self.dirs[idx], f)).convert("RGB"))
for f in files[start:start + k]
]
return torch.stack(frames) # (T, 3, H, W)
```
第三步：自监督视频预训练（掩码重建）
自监督目标选最简单的：把一部分时空 patch 遮住，让模型重建被遮住的内容。这是 MAE 系列思路在视频上的直接推广。也可以用"预测下一帧"或"对比学习"，思路相通，先用重建跑通再说。
```python
src/model.py
import torch
import torch.nn as nn
class VideoMAE(nn.Module):
"""极简版视频掩码自编码器：Conv3d 切 patch + Transformer 编码 + 反卷积重建。"""
def __init__(self, dim=384, depth=6, heads=6, patch=16, t_patch=2,
in_ch=3, max_tokens=4096, mask_ratio=0.75):
super().__init__()
self.patch = patch
self.t_patch = t_patch
self.mask_ratio = mask_ratio
self.embed = nn.Conv3d(in_ch, dim, (t_patch, patch, patch),
(t_patch, patch, patch))
self.pos = nn.Parameter(torch.zeros(1, max_tokens, dim))
self.norm = nn.AI 词典：LayerNorm">LayerNorm(dim)
layer = nn.TransformerEncoderLayer(dim, heads, dim * 4,
batch_first=True, norm_first=True)
self.encoder = nn.TransformerEncoder(layer, depth)
self.mask_token = nn.Parameter(torch.zeros(1, 1, dim))
self.decoder = nn.Sequential(
nn.ConvTranspose3d(dim, 128, (t_patch, patch, patch),
(t_patch, patch, patch)),
nn.GELU(),
nn.Conv3d(128, in_ch, 3, padding=1),
)
def features(self, x):
"""x: (B, 3, T, H, W) -> (B, dim, H/patch, W/patch)，时间维平均。"""
z = self.embed(x) # (B, dim, T', H', W')
z = z.mean(dim=2) # (B, dim, H', W')
b, d, h, w = z.shape
tokens = z.flatten(2).transpose(1, 2) + self.pos[:, :h * w]
tokens = self.encoder(self.norm(tokens))
return tokens.transpose(1, 2).reshape(b, d, h, w)
def forward(self, x):
"""x: (B, 3, T, H, W)，返回被遮住 patch 上的重建损失。"""
z = self.embed(x)
b, d, t, h, w = z.shape
tokens = z.flatten(2).transpose(1, 2) + self.pos[:, :t * h * w]
n = tokens.size(1)
keep = torch.rand(b, n, device=x.device) > self.mask_ratio
keep[:, 0] = True # 至少保留一个 token
tokens = torch.where(keep.unsqueeze(-1), tokens, self.mask_token)
feat = self.encoder(self.norm(tokens))
feat = feat.transpose(1, 2).reshape(b, d, t, h, w)
rec = self.decoder(feat)
mask = (~keep).reshape(b, 1, t, h, w).float()
loss = ((rec - x) ** 2 * mask).sum() / (mask.sum() + 1e-6)
return loss
```
训练入口，把片段的时间维和通道维顺序对齐：
```python
src/train_pretrain.py
import os
import torch
from torch.utils.data import DataLoader
from dataset import ClipDataset
from model import VideoMAE
def main():
device = "cuda" if torch.cuda.is_available() else "cpu"
os.makedirs("checkpoints", exist_ok=True)
ds = ClipDataset("data/frames", clip_len=8, size=224)
dl = DataLoader(ds, batch_size=8, shuffle=True,
num_workers=4, drop_last=True)
model = VideoMAE().to(device)
opt = torch.optim.AdamW(model.parameters(), lr=1.5e-4, weight_decay=0.05)
use_amp = device == "cuda"
scaler = torch.cuda.amp.GradScaler(enabled=use_amp)
for epoch in range(20):
model.train()
total = 0.0
for clip in dl:
clip = clip.to(device).permute(0, 2, 1, 3, 4) # (B, 3, T, H, W)
with torch.cuda.amp.autocast(enabled=use_amp):
loss = model(clip)
opt.zero_grad(set_to_none=True)
scaler.scale(loss).backward()
scaler.step(opt)
scaler.update()
total += loss.item()
print(f"epoch {epoch} loss {total / max(1, len(dl)):.4f}")
torch.save(model.state_dict(), "checkpoints/video_mae.pt")
if __name__ == "__main__":
main()
```
跑起来之后，重点看损失有没有稳定下降。如果损失很快掉到很低但下游迁移没起色，通常是任务太简单（mask 比例太低）或者数据太单一，可以调 mask 比例、加数据增强、换个数据集。
第四步：把 ARC 网格变成模型能吃的张量
ARC 的输入是 0 到 9 的整数网格，每个整数是一个颜色编号。视频编码器吃的是三通道图像，所以需要一座桥。最小可跑的桥是：给 10 个颜色配一组固定的 RGB 值，把网格当成一张小图。
注意两个关键细节：
一是必须用最近邻插值放大。双线性插值会把相邻颜色混合出中间色，模型的输入分布就完全变了。
二是网格尺寸可变。公开的 ARC 数据里网格边长最大通常不超过 30，具体以官方数据为准。统一补零到固定边长，同时额外记录一个"有效区域掩码"，否则补出来的零会被当成真实的黑色格子，训练和评测都会错。
```python
src/arc_data.py
import json
import numpy as np
import torch
import torch.nn.functional as F
一组固定调色板即可，关键是全流程保持一致
PALETTE = torch.tensor([
[0, 0, 0], [0, 116, 217], [255, 65, 54], [46, 204, 64], [255, 220, 0],
[170, 170, 170], [240, 18, 190], [255, 133, 27], [127, 219, 255], [135, 12, 37],
], dtype=torch.float32) / 255.0
GRID = 30
IMG = 224
def load_task(path):
with open(path, "r", encoding="utf-8") as f:
return json.load(f)
def grid_to_tensor(grid, size=GRID):
"""返回 (3, size, size) 图像张量和 (1, size, size) 有效区域掩码。"""
g = np.asarray(grid, dtype=np.int64)
h, w = g.shape
assert h <= size and w <= size, f"网格 {h}x{w} 超过 {size}"
rgb = PALETTE[g].permute(2, 0, 1) # (3, h, w)
x = torch.zeros(3, size, size)
valid = torch.zeros(1, size, size)
x[:, :h, :w] = rgb
valid[:, :h, :w] = 1.0
return x, valid
def to_model_input(x, img=IMG):
"""最近邻放大，保持颜色为纯色。x: (B, 3, grid, grid)"""
return F.interpolate(x, size=(img, img), mode="nearest")
```
第五步：表征迁移——用视频编码器读 ARC 网格
现在把预训练好的编码器拿过来，接一个逐格子的颜色分类头。编码器输出的是 patch 级特征，用双线性插值拉回到 30×30 的格子分辨率，每个格子做一次 10 分类。
```python
src/transfer.py
import copy
import torch
import torch.nn as nn
import torch.nn.functional as F
from arc_data import GRID, IMG, grid_to_tensor, to_model_input
class ARCPredictor(nn.Module):
def __init__(self, encoder, dim=384, n_colors=10, grid=GRID, img=IMG):
super().__init__()
self.encoder = encoder
self.grid = grid
self.img = img
self.head = nn.Linear(dim, n_colors)
def forward(self, x):
"""x: (B, 3, img, img) -> logits: (B, grid, grid, n_colors)"""
feats = self.encoder.features(x) # (B, dim, h, w)
feats = F.interpolate(feats, size=(self.grid, self.grid),
mode="bilinear", align_corners=False)
return self.head(feats.permute(0, 2, 3, 1))
def build_model(ckpt_path, device="cpu"):
from model import VideoMAE
encoder = VideoMAE()
state = torch.load(ckpt_path, map_location="cpu")
只加载编码器相关的权重，解码器在迁移阶段用不到
enc_state = {k: v for k, v in state.items()
if k.startswith(("embed", "encoder", "norm", "pos"))}
encoder.load_state_dict(enc_state, strict=False)
model = ARCPredictor(encoder).to(device)
return model
```
第六步：测试时适配
这是整条路线里最关键的一步。ARC 每个任务只给几对示例，指望一个通用模型直接答对不现实，所以要在解题时针对当前任务做适配。常用两段式：先用无标签的自监督适配调编码器，再用示例对调分类头。
自监督适配的做法是：把该任务所有输入网格（示例的输入加上测试输入）拿来做格子级掩码重建。因为颜色只有 10 种，重建本质上就是逐格分类，非常自然。
```python
src/transfer.py 追加
import numpy as np
def random_cell_mask(h, w, ratio=0.5, size=GRID):
"""在 h x w 的有效区域内随机选格子，返回 (1, size, size) 掩码。"""
m = torch.zeros(size, size)
n = max(1, int(h * w * ratio))
flat = torch.randperm(h * w)[:n]
m[flat // w, flat % w] = 1.0
return m.unsqueeze(0)
def tta_finetune(model, task, ssl_steps=100, sup_steps=200, lr=3e-4):
"""返回针对当前任务适配过的模型副本。"""
model = copy.deepcopy(model)
model.train()
冻结编码器主体，只调 LayerNorm 和分类头，样本太少时更稳
for name, p in model.encoder.named_parameters():
p.requires_grad = "norm" in name
opt = torch.optim.AdamW(
[p for p in model.parameters() if p.requires_grad], lr=lr)
inputs = [p["input"] for p in task["train"]] + [p["input"] for p in task["test"]]
阶段一：无标签的自监督适配
for step in range(ssl_steps):
grid = inputs[step % len(inputs)]
h, w = len(grid), len(grid[0])
x, valid = grid_to_tensor(grid)
m = random_cell_mask(h, w, ratio=0.5) # (1, 30, 30)
x_pert = x * (1 - m[0]) # 被遮处涂成颜色 0
logits = model(to_model_input(x_pert.unsqueeze(0))) # (1, 30, 30, 10)
target = torch.tensor(grid, dtype=torch.long)
ce = F.cross_entropy(logits[0].permute(2, 0, 1), target,
reduction="none") # (30, 30)
loss = (ce * m[0]).sum() / (m[0].sum() + 1e-6)
opt.zero_grad(set_to_none=True)
loss.backward()
opt.step()
阶段二：用示例对做监督适配
for step in range(sup_steps):
pair = task["train"][step % len(task["train"])]
x_in, valid = grid_to_tensor(pair["input"])
y_out = torch.tensor(pair["output"], dtype=torch.long)
h, w = y_out.shape
if (h, w) != (len(pair["input"]), len(pair["input"][0])):
continue # 简化版只处理同尺寸任务，见"常见坑"
logits = model(to_model_input(x_in.unsqueeze(0)))
ce = F.cross_entropy(logits[0].permute(2, 0, 1), y_out, reduction="none")
loss = (ce * valid[0]).sum() / (valid[0].sum() + 1e-6)
opt.zero_grad(set_to_none=True)
loss.backward()
opt.step()
model.eval()
return model
@torch.no_grad()
def predict_grid(model, grid):
x, _ = grid_to_tensor(grid)
logits = model(to_model_input(x.unsqueeze(0)))
pred = logits.argmax(-1)[0].cpu().numpy()
h, w = len(grid), len(grid[0])
return pred[:h, :w].tolist()
```
这里有个设计取舍：适配阶段只训练 LayerNorm 和分类头。样本只有几对，全量微调很容易过拟合到示例上，反而拉低测试表现。等你跑通之后可以再对比"只调头""调最后几层""全量微调"三种设置。
第七步：评测脚本骨架
ARC 的评测口径是整张输出网格逐格完全一致才算对，部分正确不给分。这个口径以官方评测脚本为准，自己写的时候别放宽。
```python
src/evaluate.py
import os
import glob
import torch
from arc_data import load_task
from transfer import build_model, tta_finetune, predict_grid
def main():
device = "cuda" if torch.cuda.is_available() else "cpu"
base = build_model("checkpoints/video_mae.pt", device=device)
task_files = sorted(glob.glob("data/arc/*.json"))
solved, total = 0, 0
for path in task_files:
task = load_task(path)
model = tta_finetune(base, task) # 每个任务单独适配
ok, n = 0, 0
for pair in task["test"]:
if "output" not in pair:
continue # 隐藏测试集没有答案，跳过
pred = predict_grid(model, pair["input"])
h, w = len(pair["output"]), len(pair["output"][0])
ok += int(pred == pair["output"])
n += 1
solved += ok
total += n
print(f"{os.path.basename(path)}: {ok}/{n}")
print(f"总计 {solved}/{total} = {solved / max(1, total):.3f}")
if __name__ == "__main__":
main()
```
跑完之后，你会看到哪些任务被解出来了。更值得看的是"哪些任务被解出来了"而不是总分——如果被解出来的任务集中在对称、填充、颜色映射这几类，说明视频预训练确实带进来了对应的先验。
常见坑与排错
输出尺寸不等于输入尺寸。 上面所有代码都在"输出与输入同尺寸"的简化假设下工作，而 ARC 里有相当一部分任务的输出尺寸会变。最小版本先只统计同尺寸任务，能解出来就有意义。要往前走一步，可以枚举一个候选尺寸集合，让模型额外预测输出尺寸，或者把输出尺寸当作可搜索的超参数。
颜色 0 和补零分不开。 补到 30×30 用的也是颜色 0，如果不带有效区域掩码，模型会把补出来的区域也当成任务内容去拟合。所有损失都必须乘上掩码。
放大网格用了双线性插值。 颜色会被混成新颜色，模型见到的是它从没见过的输入。网格到图像的转换里只用最近邻。
测试时适配过拟合。 症状是示例对全部答对、测试全部答错。对策：减少适配步数、降低学习率、只调 LayerNorm、加早停。也可以把示例对留一对做验证，但样本会变得更少。
数据泄漏。 测试对的输出绝不能进入适配过程。测试对的输入可以进：自监督适配本来就允许使用无标签输入，这一点在写论文式报告时要说清楚。
预训练的域差距太大。 视频是自然图像，ARC 是纯色格子，中间隔着一条沟。两个常见桥接手段：一是预训练时混入合成的网格视频（把网格的行序列当成时间轴）；二是把网格按行铺成多帧，让编码器看到"时序"结构。哪个更有效，和自己的数据量有关。
每任务单独适配很慢。 任务数一多，评测时间会成倍增长。把基础模型常驻显存、只对副本做浅层适配，或者把适配限制在几百步以内，速度会好很多。
结果不可复现。 抽帧的随机起点、掩码随机数、数据加载的多进程顺序都会影响结果。固定随机种子，并把抽帧脚本版本一起记录。
损失下降但下游没提升。 掩码比例太低、重建太容易时，编码器学到的多是低层纹理。把掩码比例调高、把片段拉长、或者换成"预测未来帧"这类更强调语义的目标，通常有帮助。
下一步建议
跑通最小版本之后，可以按这个顺序往下推：
1. 换自监督目标。掩码重建、未来帧预测、对比学习三种各跑一轮，对比下游 ARC 准确率，这是最能说明问题的一组实验。
2. 处理变尺寸输出。把尺寸预测单独拆成一个分类子任务，让主分支专心做逐格颜色预测。
3. 加任务级数据增强。颜色置换、旋转翻转、网格裁剪，都可以在适配阶段当一致性正则用，缓解示例过少的问题。
4. 做消融。随机初始化的编码器、只训分类头、全量微调三种对照跑一遍，才能说清楚"到底是谁在起作用"。
5. 用公开评测集验证。自己划分的验证集容易乐观，公开评测的提交方式与口径以官方页面说明为准。