视觉长短期记忆网络 (Vision-LSTM, ViL)

[项目主页] [论文] [模型] [代码库演示视频] [引用]

Vision-LSTM (ViL) 的 PyTorch 实现和预训练模型，这是 xLSTM 在计算机视觉领域的一种改编。

vision_lstm_schematic

许可证

本项目采用 MIT 许可证，除以下文件夹/文件采用 AGPL-3.0 许可证：

src/vislstm/modules/xlstm
vision_lstm/vision_lstm.py
vision_lstm/vision_lstm2.py

开始使用

本代码库支持 Vision-LSTM 的简单使用，包括"仅架构"实现和完整的训练流程。

仅架构

vision_lstm 包提供了一个独立的实现，风格类似于 timm。

这里提供了一个使用 vision_lstm 包在 CIFAR10 上训练 ViL 的示例。

如果您只需要模型架构，可以通过 torchhub 单行加载，或将 vision_lstm 文件夹复制到您自己的代码库中。请注意，对于 VisionLSTM2，我们将一个块视为由两个子块组成（第一个从右上到左下，第二个从右下到左上），以便更容易实现分层学习率衰减。

# 加载 ViL-T
model = torch.hub.load("nx-ai/vision-lstm", "VisionLSTM2")
# 加载您自己的模型
model = torch.hub.load(
    "nx-ai/vision-lstm", 
    "VisionLSTM2",  # VisionLSTM2 是 VisionLSTM 的改进版本
    dim=192,  # 潜在维度（ViL-T 为 192）
    depth=12,  # ViL 块的数量（1 个块由前向和后向块的 2 个子块组成）
    patch_size=16,  # 补丁大小（224x224 图像会产生 196 个补丁）
    input_shape=(3, 224, 224),  # RGB 图像，分辨率为 224x224
    output_shape=(1000,),  # 1000 类分类器
    drop_path_rate=0.05,  # 随机深度参数
)

有关 VisionLSTM 和 VisionLSTM2 之间的变化列表，请参见下文或附录 A。我们建议使用 VisionLSTM2，因为我们发现它表现更好，但保留 VisionLSTM 以保持向后兼容性。

完整的训练/评估流程（架构、数据集、超参数、分类、分割等）

如果您想用我们的代码库训练模型，请按照 SETUP.md 中的设置说明进行操作。要开始运行，请按照 RUN.md 中的说明进行操作。

所有训练/评估运行的配置/超参数可以在这里找到。

VTAB-1K 评估使用这个代码库进行。

预训练模型

在 ImageNet-1K 上预训练的模型可以通过 torchhub 加载，或直接从这里下载。

# ImageNet-1K 预训练模型
model = torch.hub.load("nx-ai/vision-lstm", "vil2-tiny")               # 78.3%
model = torch.hub.load("nx-ai/vision-lstm", "vil2-small")              # 81.5%
model = torch.hub.load("nx-ai/vision-lstm", "vil2-base")               # 82.4%

# 仅训练 400 轮的 ViL-T（附录 B.2）
model = torch.hub.load("nx-ai/vision-lstm", "vil2-tiny-e400")          # 77.2%

这些模型的预训练日志可以在这里找到。

如何使用这些模型的示例可以在 eval.py 中找到，该文件在 ImageNet-1K 验证集上评估这些模型。

DeiT-III-T 重新实现模型

我们重新实现的 DeiT-III-T 的检查点可以在这里下载原始检查点，也可以从 torchhub 加载（视觉变换器实现基于 KappaModules，所以在通过 torchhub 加载 ViT 检查点之前，你需要运行 pip install kappamodules==0.1.70 来安装它）。

model = torch.hub.load("nx-ai/vision-lstm", "deit3-tiny-e400")  # 75.6%
model = torch.hub.load("nx-ai/vision-lstm", "deit3-tiny")       # 76.2%

版本1预训练模型

在 ViL 的第一次迭代中，模型训练时使用了以下设置：(i) 双边平均池化而不是双边拼接 (ii) 在 q 和 k 之前使用因果卷积1d而不是卷积2d (iii) 投影和层归一化中没有偏置 (iv) 整个训练过程使用224分辨率，而不是先在192分辨率上预训练然后在224分辨率上短暂微调。这些改变将 ViL-T 在 ImageNet-1K 上的准确率从77.3%提高到78.3%。更多细节请参见论文附录A。我们建议使用 VisionLSTM2 而不是 VisionLSTM，但仍然保持对初始版本的原样支持。第一次迭代的预训练模型可以按以下方式加载：

# ImageNet-1K 预训练模型
model = torch.hub.load("nx-ai/vision-lstm", "vil-tiny")               # 77.3%
model = torch.hub.load("nx-ai/vision-lstm", "vil-tinyplus")           # 78.1%
model = torch.hub.load("nx-ai/vision-lstm", "vil-small")              # 80.7%
model = torch.hub.load("nx-ai/vision-lstm", "vil-smallplus")          # 80.9%
model = torch.hub.load("nx-ai/vision-lstm", "vil-base")               # 81.6%

# 长序列微调模型
model = torch.hub.load("nx-ai/vision-lstm", "vil-tinyplus-stride8")   # 80.0%
model = torch.hub.load("nx-ai/vision-lstm", "vil-smallplus-stride8")  # 82.2%
model = torch.hub.load("nx-ai/vision-lstm", "vil-base-stride8")       # 82.7%

# 仅训练400轮的tiny模型
model = torch.hub.load("nx-ai/vision-lstm", "vil-tiny-e400")          # 76.1%
model = torch.hub.load("nx-ai/vision-lstm", "vil-tinyplus-e400")      # 77.2%

使用随机权重初始化可以按以下方式进行：

# 加载 ViL-T
model = torch.hub.load("nx-ai/vision-lstm", "VisionLSTM")
# 加载你自己的模型
model = torch.hub.load(
    "nx-ai/vision-lstm", 
    "VisionLSTM",
    dim=192,  # 潜在维度（ViL-T 为192）
    depth=24,  # ViL 块的数量
    patch_size=16,  # 图像块大小（对于224x224的图像会产生196个图像块）
    input_shape=(3, 224, 224),  # RGB图像，分辨率为224x224
    output_shape=(1000,),  # 1000类的分类器
    drop_path_rate=0.05,  # 随机深度参数
    stride=None,  # 设置为8用于长序列微调
)

其他

这个代码库是 MIM-Refiner 使用的代码库的改进版本，有一个演示视频解释了各种细节。

VTAB-1K 评估是使用这个代码库进行的。

引用

如果你喜欢我们的工作，请考虑给它一个星标 :star: 并引用我们

@article{alkin2024visionlstm,
  title={Vision-LSTM: xLSTM as Generic Vision Backbone},
  author={Benedikt Alkin and Maximilian Beck and Korbinian P{\"o}ppel and Sepp Hochreiter and Johannes Brandstetter}
  journal={arXiv preprint arXiv:2406.04303},
  year={2024}
}