我们正式发布 TorchMultimodal Beta,这是一个用于大规模训练最先进(SoTA)多任务多模态模型的 PyTorch 领域库。该库提供了可组合的构建模块(模块、变换、损失函数)以加速模型开发,包含来自已发表研究的 SoTA 模型架构(FLAVA、MDETR、Omnivore)、训练和评估脚本,以及用于探索这些模型的 Notebook。该库正处于积极开发阶段,我们非常期待听到您的反馈!您可以在此处了解更多入门详细信息。
为什么要选择 TorchMultimodal?
人们对理解多种输入类型(文本、图像、视频和音频信号)的人工智能模型越来越感兴趣,并希望利用这种理解来生成不同形式的输出(句子、图片、视频)。来自 FAIR 的近期研究,如 FLAVA、Omnivore 和 data2vec,已经表明用于理解的多模态模型在性能上已可与单模态模型媲美,在某些情况下甚至确立了新的最先进水平。诸如 Make-a-video 和 Make-a-scene 等生成式模型正在重新定义现代人工智能系统的能力。
随着人们对多模态 AI 兴趣的增长,研究人员正在寻找能够快速验证想法并基于领域最新研究成果进行构建的工具和库。虽然 PyTorch 生态系统拥有丰富的库和框架,但这些组件如何互操作,或者如何将它们拼接在一起以构建 SoTA 多模态模型,往往并不直观。
TorchMultimodal 通过提供以下内容解决了这一问题:
- 可组合且易于使用的构建模块:研究人员可以使用这些模块来加速各自工作流程中的模型开发和实验。它们被设计为模块化,并且可以轻松扩展以处理新的模态。
- 用于训练和评估最新研究模型的端到端示例:这些示例既可作为正在进行或未来研究的起点,也可作为使用高级功能(如集成 FSDP 和激活检查点以扩展模型和批处理规模)的示范。
介绍 TorchMultimodal
TorchMultimodal 是一个用于大规模训练多任务多模态模型的 PyTorch 领域库。在存储库中,我们提供:
- 构建模块:一系列模块化且可组合的构建块,如模型、融合层、损失函数、数据集和实用工具。一些示例包括:
- 带温度系数的对比损失(Contrastive Loss with Temperature):用于训练 CLIP 和 FLAVA 等模型的常用函数。我们还包括了诸如 ImageTextContrastiveLoss(用于 ALBEF 等模型)等变体。
- 码本(Codebook)层:通过在嵌入空间中进行最近邻查找来压缩高维数据,是 VQVAE 的重要组成部分(在存储库中作为模型提供)。
- 移动窗口注意力机制(Shifted-window Attention):基于窗口的多头自注意力机制,是 Swin 3D Transformer 等编码器的重要组成部分。
- CLIP 组件:由 OpenAI 发布的一款流行模型,已被证明在学习文本和图像表示方面极其有效。
- 多模态 GPT:一种抽象模型,当与生成实用工具结合使用时,可将 OpenAI 的 GPT 架构扩展到多模态生成任务。
- 多头注意力机制(MultiHeadAttention):基于注意力模型的关键组件,支持快速自回归解码。
- 示例:一系列展示如何将这些构建块与 PyTorch 生态系统中(如 Lightning、TorchMetrics)的组件和通用基础设施相结合,以复现文献中发表的最先进模型的示例。我们目前提供五个示例,包括:
- FLAVA [论文]:被 CVPR 录用的官方论文代码,包含微调 FLAVA 的教程。
- MDETR [论文]:与纽约大学(NYU)作者合作提供的示例,解决了 PyTorch 生态系统中的互操作性痛点,包括一个关于使用 MDETR 进行短语定位和视觉问答的 Notebook。
- Omnivore [论文]:TorchMultimodal 中首个处理视频和 3D 数据的模型示例,包含一个用于探索该模型的 Notebook。
- MUGEN [论文]:用于自回归生成和检索的奠基性工作,包括使用从 OpenAI coinrun 扩充的大规模合成数据集进行文本生成视频和检索的演示。
- ALBEF [论文]:该模型的代码,包括一个用于利用该模型进行视觉问答的 Notebook。
以下代码片段展示了与 CLIP 相关的多个 TorchMultimodal 组件的用法示例
# instantiate clip transform
clip_transform = CLIPTransform()
# pass the transform to your dataset. Here we use coco captions
dataset = CocoCaptions(root= ..., annFile=..., transforms=clip_transform)
dataloader = DataLoader(dataset, batch_size=16)
# instantiate model. Here we use clip with vit-L as the image encoder
model= clip_vit_l14()
# define loss and other things needed for training
clip_loss = ContrastiveLossWithTemperature()
optim = torch.optim.AdamW(model.parameters(), lr = 1e-5)
epochs = 1
# write your train loop
for _ in range(epochs):
for batch_idx, batch in enumerate(dataloader):
image, text = batch
image_embeddings, text_embeddings = model(image, text)
loss = contrastive_loss_with_temperature(image_embeddings, text_embeddings)
loss.backward()
optimizer.step()
除了代码之外,我们还发布了一份多模态基础模型微调教程,以及一篇关于如何使用 PyTorch Distributed 技术(FSDP 和激活检查点)扩展此类模型的博客文章(附代码指向)。我们希望这些示例和教程能够揭开 PyTorch 生态系统中可用的一系列高级功能的神秘面纱。
下一步计划
虽然这是一个激动人心的发布,但未来还有更多内容。该库正在开发中,我们正致力于加入扩散模型领域的一些令人兴奋的发展成果,以及展示研究中常见趋势的示例。在您探索和使用该库时,我们非常乐意听到您的任何反馈!您可以在此处了解更多入门详细信息。
团队
TorchMultimodal 的主要贡献者和开发者包括 Ankita De、Evan Smothers、Kartikay Khandelwal、Lan Gong、Laurence Rouesnel、Nahiyan Malik、Rafi Ayub 和 Yosua Michael Maranatha。