由 Kartikay Khandelwal,Ankita De 撰写

我们宣布推出 TorchMultimodal Beta,这是一个 PyTorch 领域库,用于大规模训练最先进的任务多模态模型。该库提供了可组合的构建块(模块、转换、损失函数),以加速模型开发,包括从已发表研究中提取的最先进的模型架构(FLAVA、MDETR、Omnivore),以及用于探索这些模型的训练和评估脚本以及笔记本。该库正在积极开发中,我们非常欢迎您的反馈!您可以通过此处了解更多如何开始使用的信息。

为什么选择 TorchMultimodal?

对理解多种输入类型(文本、图像、视频和音频信号)的 AI 模型感兴趣,并且可以选择使用这种理解来生成不同形式的输出(句子、图片、视频)。FAIR 的最新工作,如 FLAVA、Omnivore 和 data2vec,表明理解的多模态模型与单模态模型具有竞争力,在某些情况下甚至确立了新的最先进水平。生成模型如 Make-a-video 和 Make-a-scene 正在重新定义现代 AI 系统可以做什么。

随着对多模态 AI 的兴趣日益增长,研究人员正在寻找工具和库来快速实验想法,并在此基础上构建该领域的最新研究成果。虽然 PyTorch 生态系统拥有丰富的库和框架,但组件之间的互操作性或如何将这些组件拼接在一起以构建最先进的多模态模型并不总是显而易见。

TorchMultimodal 通过提供以下功能解决了这个问题:

  • 可组合且易于使用的构建块,研究人员可以使用这些构建块加速他们在自己工作流程中的模型开发和实验。这些构建块设计为模块化,可以轻松扩展以处理新的模态。

  • 训练和评估最新研究模型的端到端示例。这些示例应作为持续/未来研究的起点,以及使用高级功能(如与 FSDP 集成和激活检查点以扩展模型和批量大小)的示例。

介绍 TorchMultimodal

TorchMultimodal 是一个用于大规模训练多任务多模态模型的 PyTorch 领域库。在仓库中,我们提供了:

  • 构建模块。一组模块化和可组合的构建模块,如模型、融合层、损失函数、数据集和实用工具。一些示例包括:

    • 带温度的对比损失。常用于训练 CLIP 和 FLAVA 等模型。我们还包括变体,如用于 ALBEF 等模型的 ImageTextContrastiveLoss。

    • 压缩高维数据的代码簿层,通过嵌入空间中的最近邻查找实现,是 VQVAEs(在仓库中提供作为模型)的一个关键组件。

    • 基于移位窗口的多头自注意力机制,是 Swin 3D Transformer 等编码器的一个关键组件。

    • CLIP 的组件。由 OpenAI 发布的一个流行模型,已被证明在学习和文本及图像表示方面非常有效。

    • 多模态 GPT。一种抽象,它扩展了 OpenAI 的 GPT 架构,用于与生成工具结合时的多模态生成。

    • 多头注意力机制。支持快速自回归解码的基于注意力的模型的关键组件。

  • 示例。一组示例,展示了如何将这些构建块与 PyTorch 生态系统中的组件和常见基础设施(Lightning、TorchMetrics)结合使用,以复制文献中发布的最先进模型。我们目前提供五个示例,包括。

    • FLAVA [论文]。CVPR 会议接受论文的官方代码,包括 FLAVA 微调教程。

    • MDETR [论文]。与 NYU 作者合作,提供一个示例,该示例缓解了 PyTorch 生态系统中的互操作性痛点,包括使用 MDETR 进行短语定位和视觉问答的笔记本。

    • Omnivore [论文]。TorchMultimodal 中处理视频和 3D 数据的第一个示例,包括用于探索该模型的笔记本。

    • MUGEN [论文]。自回归生成和检索的基础性工作,包括使用来自 OpenAI coinrun 的大规模合成数据集进行文本-视频生成和检索的演示。

    • ALBEF [论文] 模型代码,包括用于视觉问答的该模型使用笔记本。

以下代码片段展示了几个与 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 分布式(FSDP 和激活检查点技术)来扩展此类模型。我们希望这些示例和教程能够帮助大家揭开 PyTorch 生态系统中的许多高级功能的神秘面纱。

接下来是什么?

虽然这是一个令人兴奋的发布,但还有更多内容即将到来。这个库仍在开发中,我们正在努力添加扩散模型领域的激动人心的进展以及展示研究中的常见趋势的示例。当您探索和使用这个库时,我们非常乐意听取您的任何反馈!有关如何开始的更多详细信息,请在此处查看。

团队

TorchMultimodal 的主要贡献者和开发者包括 Ankita De、Evan Smothers、Kartikay Khandelwal、Lan Gong、Laurence Rouesnel、Nahiyan Malik、Rafi Ayub 和 Yosua Michael Maranatha。