lucidrains/rectified-flow-pytorch

rectified-flow-pytorch – Rectified FlowのPyTorch実装

概要 – Rectified Flow生成モデル技術（Liuら、2022年）とその後の改良を再現する、pipインストール可能な小さなライブラリです。すぐに使えるRectifiedFlowクラス、シンプルなUNetバックボーン、画像データセットを読み込むためのユーティリティ、そして🤗 Accelerateに基づくトレーナーを提供します。

重要性 – Rectified Flowは、フローマッチング/拡散に似た手法で、ノイズからデータへの直線的なマッピングを1ステップで学習し、高速なサンプリングと競争力のある画質を実現します。このリポジトリにより、研究者や愛好家は、数式をゼロから書くことなく、この手法を実験できます。

---

クイックスタート

pip install rectified-flow-pytorch          # コアライブラリ
pip install .[examples]                     # オプションのサンプル依存関係

import torch
from rectified_flow_pytorch import RectifiedFlow, Unet

# 小さなUNetバックボーン（任意のtorch.nn.Moduleに置き換え可能）
model = Unet(dim=64)

# RectifiedFlowラッパーで包む
flow = RectifiedFlow(model)

# ダミーデータ – 形状 (バッチ, 3, 256, 256)
imgs = torch.randn(1, 3, 256, 256)

# フォワードパスはバックプロパゲーション可能な損失を返す
loss = flow(imgs)
loss.backward()

# サンプルを生成（入力画像と同じ形状）
sample = flow.sample()
assert sample.shape[1:] == imgs.shape[1:]

リフロー（反復的改良）

ライブラリには、元の論文で説明されているように、学習したフローを繰り返し適用できるReflowラッパーも付属しています。

from rectified_flow_pytorch import Reflow

reflow = Reflow(flow)
reflow_loss = reflow()
reflow_loss.backward()

sample = reflow.sample()

自分の画像でトレーニング

from rectified_flow_pytorch import ImageDataset, Trainer

train_set = ImageDataset(folder='./my_images', image_size=256)
trainer = Trainer(
    flow,
    dataset=train_set,
    num_train_steps=70_000,
    results_folder='./results'   # 定期的にサンプルを保存
)
trainer()   # トレーニングループを開始（内部でaccelerateを使用）

---

内部構成

コンポーネント	役割
RectifiedFlow	直線フロー目的関数を実装するコアの損失とサンプラー。
Reflow	フローを複数回適用するための便利なラッパー（「リフロー」手順）。
Unet	デフォルトのバックボーンとして使用される最小限のUNetアーキテクチャ。dim引数で完全に設定可能。
ImageDataset	フォルダからすべての画像を読み込み、正方形サイズにリサイズするシンプルなtorch.utils.data.Dataset。
Trainer	🤗 Accelerateに基づくトレーニングループ。オプティマイザ、学習率スケジュール、チェックポイント、定期的な画像保存を処理。

---

サンプルスクリプト

リポジトリには、Oxford-Flowersデータセット（バッチ32、11,000ステップ）でモデルをトレーニングするtrain_oxford.pyが含まれています。オプションのexamples依存関係をインストール後、以下を実行できます：

python train_oxford.py

このスクリプトは、データセットの読み込み、トレーナー設定、生成サンプルの可視化方法を示します。

---

参考文献と関連研究

READMEには、元のRectified Flow論文と、フローマッチング手法を拡張または分析する新しいarXivの作品（例：Improving the Training of Rectified Flows、Scaling Rectified Flow Transformers、Consistency Flow Matching）がリストされています。これらの引用は、ユーザーが学術的な文脈を見つけ、高度なバリアントを探求するのに役立ちます。

---

誰が使うべきか？

• 研究者 – フローマッチング、拡散の代替、またはワンステップ生成モデルを実験している方。
• 実務者 – 高速サンプリング画像生成器が必要で、軽量なPyTorch実装を求める方。
• 学生 – Rectified Flowの数学を学ぶための簡潔なコードベースを探している方。

---

結論

rectified-flow-pytorchは、最先端の生成モデル技術の焦点を絞った、十分に文書化された実装です。単一のpip installで、使用可能なモデル、トレーニングユーティリティ、サンプルスクリプトを入手でき、独自の画像データでRectified Flowを簡単にプロトタイプ化およびベンチマークできます。