我们在此宣布 TorchMultimodal Beta 版的发布,这是一个 PyTorch 领域库,用于大规模训练先进的多任务多模态模型。该库提供可组合的构建块(模块、转换、损失函数)以加速模型开发,提供了已发表研究中的先进模型架构(FLAVA、MDETR、Omnivore),以及用于探索这些模型的训练、评估脚本和 notebook。该库正在积极开发中,我们期待您的反馈!您可以在此处找到有关如何上手的更多详细信息。
为何选择 TorchMultimodal?
人们对能够理解多种输入类型(文本、图像、视频和音频信号)并可选择性地利用这种理解来生成不同形式输出(句子、图片、视频)的人工智能模型越来越感兴趣。FAIR 最近的研究成果,如 FLAVA、Omnivore 和 data2vec,已经证明了用于理解的多模态模型与单模态模型相比具有竞争力,在某些情况下甚至树立了新的技术标杆。而诸如 Make-a-video 和 Make-a-scene 等生成模型正在重新定义现代人工智能系统的能力。
随着对多模态人工智能兴趣的增长,研究人员正在寻找能够快速试验想法并基于该领域最新研究成果进行构建的工具和库。虽然 PyTorch 生态系统拥有丰富的库和框架,但这些组件之间如何互操作,或者如何将它们组合起来构建先进的多模态模型,并不总是显而易见的。
TorchMultimodal 通过提供以下内容解决了这个问题:
- 可组合且易于使用的构建块,研究人员可以使用它们来加速自己工作流程中的模型开发和实验。这些构建块被设计为模块化的,并且可以轻松扩展以处理新的模态。
- 用于训练和评估最新研究模型的端到端示例。这些示例应作为当前/未来研究的起点,并作为使用高级功能的范例,例如与 FSDP 集成以及使用激活检查点来扩展模型和批量大小。
TorchMultimodal 简介
TorchMultimodal 是一个 PyTorch 领域库,用于大规模训练多任务多模态模型。在该代码库中,我们提供:
- 构建块。一系列模块化和可组合的构建块,如模型、融合层、损失函数、数据集和实用工具。一些示例包括:
- 带温度系数的对比损失。常用于训练 CLIP 和 FLAVA 等模型的函数。我们还包括了诸如 ALBEF 等模型中使用的变体,如 ImageTextContrastiveLoss。
- 码本层,通过在嵌入空间中进行最近邻查找来压缩高维数据,是 VQVAE(在代码库中作为模型提供)的重要组成部分。
- 位移窗口注意力,基于窗口的多头自注意力机制,是 Swin 3D Transformers 等编码器的重要组成部分。
- CLIP 的组件。由 OpenAI 发布的一款流行模型,已证明在学习文本和图像表示方面非常有效。
- 多模态 GPT。一个抽象模型,当与生成实用工具结合时,可将 OpenAI 的 GPT 架构扩展用于多模态生成。
- MultiHeadAttention。基于注意力模型的一个关键组件,支持快速自回归解码。
- 示例。一系列示例,展示了如何将这些构建块与 PyTorch 生态系统中的组件和通用基础架构(Lightning、TorchMetrics)相结合,以复现文献中发布的先进模型。我们目前提供五个示例,包括:
- FLAVA [论文]。被 CVPR 接收的论文的官方代码,包括一个关于微调 FLAVA 的教程。
- MDETR [论文]。与纽约大学的作者合作提供的示例,旨在解决 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。