Perceiver、Perceiver IO 和 Perceiver AR

本仓库是 Perceiver、Perceiver IO 和 Perceiver AR 的 PyTorch 实现，并提供了用于模型训练的 PyTorch Lightning 接口以及用于推理的 Hugging Face 🤗 接口。

Perceiver: 基于迭代注意力机制的通用感知模型 (论文, 视频)
Perceiver IO: 面向结构化输入与输出的通用架构 (论文, 博客文章)
使用 Perceiver AR 进行通用的长上下文自回归建模 (论文, 博客文章)

概述

perceiver-io 库的核心是 后端模型，即 Perceiver、Perceiver IO 和 Perceiver AR 的轻量级 PyTorch 实现。这些模型可以被封装为 PyTorch Lightning 模块以进行训练（Lightning 接口），也可以被封装为 🤗 模块以用于推理（Hugging Face 接口）。详细信息请参阅 库设计。

训练的命令行接口使用 Lightning CLI 实现。训练数据集是 🤗 datasets，并被封装为 PyTorch Lightning 数据模块。对于 NLP 任务，perceiver-io 支持所有 🤗 fast tokenizers 以及 🤗 Perceiver UTF-8 字节分词器。

文档

• 安装
• 快速入门
• 库设计
• 预训练模型
• 训练示例
• 推理示例
• 模型构建
• 基础组件

安装

通过 pip

pip install perceiver-io[text,vision,audio]

从源码安装

从源码安装需要先安装 Miniconda 和 Poetry（版本 1.2.0 或更高）。

创建并激活 perceiver-io 的 conda 环境：

conda env create -f environment.yml
conda activate perceiver-io

安装主依赖和测试依赖，包括所有额外依赖：

# 不包含示例所需的依赖
poetry install --all-extras

如果希望在本地运行 示例，则需额外添加 --with examples：

poetry install --all-extras --with examples

Docker 镜像

docker pull ghcr.io/krasserm/perceiver-io:latest

详细信息请参阅 Docker 镜像。

快速入门

推理

光流

计算输入视频中连续帧之间的光流，并将渲染结果写入输出视频：

from urllib.request import urlretrieve
from transformers import pipeline

from perceiver.data.vision import video_utils
from perceiver.model.vision import optical_flow  # 注册自动类和流水线

urlretrieve(
    url="https://martin-krasser.com/perceiver/flow/sintel_clip_cave_dragon_fight.mp4",
    filename="sintel_clip_cave_dragon_fight.mp4",
)

# 创建光流流水线
optical_flow_pipeline = pipeline("optical-flow", model="krasserm/perceiver-io-optical-flow", device="cuda:0")

# 加载连续的视频帧对
frame_pairs = video_utils.read_video_frame_pairs("sintel_clip_cave_dragon_fight.mp4")

# 为所有帧对生成并渲染光流
optical_flows = optical_flow_pipeline(frame_pairs, render=True, device="cuda:0")

# 创建包含渲染光流的视频
video_utils.write_video("sintel_clip_cave_dragon_fight_output.mp4", optical_flows, fps=24)

以下是输入视频与输出视频的并排对比：

符号音频生成

通过生成符号化（MIDI)音频数据，并使用 fluidsynth 将生成的音频符号转换为 WAV 输出来创建音频序列（注意： 必须先安装 fluidsynth 才能使以下示例正常工作）：

from transformers import pipeline
from pretty_midi import PrettyMIDI
from perceiver.model.audio import symbolic  # 自动类注册

repo_id = "krasserm/perceiver-ar-sam-giant-midi"

prompt = PrettyMIDI("prompt.mid")
audio_generator = pipeline("symbolic-audio-generation", model=repo_id)

output = audio_generator(prompt, max_new_tokens=64, num_latents=1, do_sample=True, top_p=0.95, temperature=1.0, render=True)

with open("generated_audio.wav", "wb") as f:
    f.write(output["generated_audio_wav"])

生成的音频序列示例可在 🤗 hub 上找到。

更多示例请参阅 推理示例。

训练

在命令行上训练一个小型 Perceiver IO 图像分类器（907K 参数），数据集为 MNIST。该分类器通过 重复的交叉注意力 对输入图像的单个像素进行跨注意力操作。更多详细信息请参阅 图像分类 训练示例。

python -m perceiver.scripts.vision.image_classifier fit \
  --model.num_latents=32 \
  --model.num_latent_channels=128 \
  --model.encoder.num_frequency_bands=32 \
  --model.encoder.num_cross_attention_layers=2 \
  --model.encoder.num_self_attention_blocks=3 \
  --model.encoder.num_self_attention_layers_per_block=3 \
  --model.encoder.first_self_attention_block_shared=false \
  --model.encoder.dropout=0.1 \
  --model.encoder.init_scale=0.1 \
  --model.decoder.num_output_query_channels=128 \
  --model.decoder.dropout=0.1 \
  --model.decoder.init_scale=0.1 \
  --data=MNISTDataModule \
  --data.batch_size=64 \
  --optimizer=AdamW \
  --optimizer.lr=1e-3 \
  --lr_scheduler.warmup_steps=500 \
  --trainer.accelerator=gpu \
  --trainer.devices=1 \
  --trainer.max_epochs=30 \
  --trainer.logger=TensorBoardLogger \
  --trainer.logger.save_dir=logs \
  --trainer.logger.name=logs

模型构建 描述了如何使用 Lightning CLI 实现特定于模型的命令行界面。训练检查点将被写入 logs/img_clf/version_0/checkpoints 目录。假设存在一个文件名为 epoch=025-val_loss=0.065.ckpt 的检查点，可以通过以下方式将其转换为 perceiver-io 🤗 模型：

from perceiver.model.vision.image_classifier import convert_mnist_classifier_checkpoint

convert_mnist_classifier_checkpoint(
    save_dir="example/mnist-classifier",
    ckpt_url="logs/img_clf/version_0/checkpoints/epoch=025-val_loss=0.065.ckpt",
)

这样就可以将其用于 🤗 图像分类流水线中：

from datasets import load_dataset
from transformers import pipeline

mnist_dataset = load_dataset("mnist", split="test")[:9]

images = mnist_dataset["image"]
labels = mnist_dataset["label"]

classifier = pipeline("image-classification", model="example/mnist-classifier")
predictions = [pred[0]["label"] for pred in classifier(images)]

print(f"Labels:      {labels}")
print(f"Predictions: {predictions}")

Labels:      [7, 2, 1, 0, 4, 1, 4, 9, 5]
Predictions: [7, 2, 1, 0, 4, 1, 4, 9, 5]

或者直接加载：

import torch
from transformers import AutoModelForImageClassification, AutoImageProcessor

model = AutoModelForImageClassification.from_pretrained("example/mnist-classifier")
processor = AutoImageProcessor.from_pretrained("example/mnist-classifier")

inputs = processor(images, return_tensors="pt")

with torch.no_grad():
    # 使用 perceiver-io Hugging Face 模型
    output_1 = model(**inputs).logits

with torch.no_grad():
    # 或者直接使用 perceiver-io 后端模型
    output_2 = model.backend_model(inputs.pixel_values)

print(f"Predictions: {output_1.argmax(dim=-1).numpy().tolist()}")
print(f"Predictions: {output_2.argmax(dim=-1).numpy().tolist()}")

Predictions: [7, 2, 1, 0, 4, 1, 4, 9, 5]
Predictions: [7, 2, 1, 0, 4, 1, 4, 9, 5]

更多示例请参阅 训练示例。

文章

引用本仓库的文章：

• 训练计算最优的 Perceiver AR 语言模型
• 旋转位置编码的入门介绍

其他实现

• Perceiver
• Perceiver IO
• Perceiver AR

Perceiver IO 快速上手指南

Perceiver IO 是一个基于 PyTorch 的通用架构库，实现了 Perceiver、Perceiver IO 和 Perceiver AR 模型。它支持处理结构化输入与输出（如图像、音频、视频等），并提供 PyTorch Lightning 训练接口和 Hugging Face 🤗 推理接口。

环境准备

系统要求

• 操作系统: Linux, macOS, Windows (WSL 推荐)
• Python: 3.8+
• GPU: 可选，但推理和训练大型模型时强烈建议使用 CUDA 兼容的 NVIDIA GPU

前置依赖

若选择从源码安装，需预先安装以下工具：

• Miniconda
• Poetry (版本 1.2.0 或更高)

提示：国内用户可使用清华或阿里镜像源加速 Conda 和 Pip 下载。

• Conda 配置：conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main/
• Pip 配置：pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

方式一：通过 Pip 安装（推荐）

直接安装核心库及常用扩展（文本、视觉、音频）：

pip install perceiver-io[text,vision,audio]

方式二：从源码安装

适用于需要运行示例代码或进行开发的场景。

1. 创建并激活 Conda 环境：

conda env create -f environment.yml
conda activate perceiver-io

2. 安装依赖（包含所有额外功能）：

poetry install --all-extras

3. 若需本地运行示例代码，额外执行：

poetry install --all-extras --with examples

方式三：Docker

直接使用预构建的镜像：

docker pull ghcr.io/krasserm/perceiver-io:latest

基本使用

Perceiver IO 的核心优势在于其统一的推理接口。以下展示两个最典型的用法：视觉任务（光流估计）和音频生成。

1. 视觉任务：视频光流估计

计算连续视频帧之间的光流并生成可视化视频。

from urllib.request import urlretrieve
from transformers import pipeline

from perceiver.data.vision import video_utils
from perceiver.model.vision import optical_flow  # register auto-classes and pipeline

# 下载示例视频
urlretrieve(
    url="https://martin-krasser.com/perceiver/flow/sintel_clip_cave_dragon_fight.mp4",
    filename="sintel_clip_cave_dragon_fight.mp4",
)

# 创建光流推理管道
optical_flow_pipeline = pipeline("optical-flow", model="krasserm/perceiver-io-optical-flow", device="cuda:0")

# 读取连续视频帧对
frame_pairs = video_utils.read_video_frame_pairs("sintel_clip_cave_dragon_fight.mp4")

# 生成并渲染光流
optical_flows = optical_flow_pipeline(frame_pairs, render=True, device="cuda:0")

# 输出结果视频
video_utils.write_video("sintel_clip_cave_dragon_fight_output.mp4", optical_flows, fps=24)

2. 音频任务：符号化音频生成

基于 MIDI 提示生成新的音频序列，并转换为 WAV 文件。 (注意：此功能需系统预先安装 fluidsynth)

from transformers import pipeline
from pretty_midi import PrettyMIDI
from perceiver.model.audio import symbolic  # auto-class registration

repo_id = "krasserm/perceiver-ar-sam-giant-midi"

# 加载提示音
prompt = PrettyMIDI("prompt.mid")
audio_generator = pipeline("symbolic-audio-generation", model=repo_id)

# 生成音频
output = audio_generator(prompt, max_new_tokens=64, num_latents=1, do_sample=True, top_p=0.95, temperature=1.0, render=True)

# 保存为 WAV 文件
with open("generated_audio.wav", "wb") as f:
    f.write(output["generated_audio_wav"])

3. 自定义模型推理

如果您有自己训练或转换的模型（例如 MNIST 分类器），可以直接使用 Hugging Face 的标准接口：

from transformers import pipeline

# 加载本地或 Hub 上的模型
classifier = pipeline("image-classification", model="example/mnist-classifier")

# 推理
predictions = classifier(images)