TorchMultimodal (测试版发布)

模型 | 示例脚本 | 入门 | 代码概览 | 安装 | 贡献 | 许可证

简介

TorchMultimodal是一个PyTorch库，用于大规模训练最先进的多模态多任务模型，包括内容理解和生成模型。TorchMultimodal包含：

• 模块化和可组合的构建块仓库（融合层、损失函数、数据集和工具）。
• 由上述构建块组成的常见多模态模型类集合，并为规范配置提供预训练权重。
• 一系列示例，展示如何将这些构建块与PyTorch生态系统中的组件和常见基础设施结合，以复现文献中发表的最先进模型。这些示例可作为该领域持续研究的基准，以及未来工作的起点。

模型

TorchMultimodal包含多个模型，包括

• ALBEF：模型类，论文
• BLIP-2：模型类，论文
• CLIP：模型类，论文
• CoCa：模型类，论文
• DALL-E 2：模型，论文
• FLAVA：模型类，论文
• MAE/Audio MAE：模型类，MAE论文，Audio MAE论文
• MDETR：模型类，论文

示例脚本

除了上述模型外，我们还提供了用于训练、微调和评估流行多模态任务模型的示例脚本。示例可在examples/下找到，包括

模型	支持的任务
ALBEF	检索 <br/> 视觉问答
DDPM	训练和推理 (笔记本)
FLAVA	预训练 <br/> 微调 <br/> 零样本
MDETR	短语定位 <br/> 视觉问答
MUGEN	文本到视频检索 <br/> 文本到视频生成
Omnivore	预训练 <br/> 评估

入门

以下我们给出如何使用TorchMultimodal的组件编写简单训练或零样本评估脚本的最小示例。

<details> <summary>FLAVA零样本示例</summary>

import torch
from PIL import Image
from torchmultimodal.models.flava.model import flava_model
from torchmultimodal.transforms.bert_text_transform import BertTextTransform
from torchmultimodal.transforms.flava_transform import FLAVAImageTransform

# 定义零样本预测的辅助函数
def predict(zero_shot_model, image, labels):
  zero_shot_model.eval()
  with torch.no_grad():
      image = image_transform(img)["image"].unsqueeze(0)
      texts = text_transform(labels)
      _, image_features = zero_shot_model.encode_image(image, projection=True)
      _, text_features = zero_shot_model.encode_text(texts, projection=True)
      scores = image_features @ text_features.t()
      probs = torch.nn.Softmax(dim=-1)(scores)
      label = labels[torch.argmax(probs)]
      print(
          "标签概率: ",
          {labels[i]: probs[:, i] for i in range(len(labels))},
      )
      print(f"预测标签: {label}")


image_transform = FLAVAImageTransform(is_train=False)
text_transform = BertTextTransform()
zero_shot_model = flava_model(pretrained=True)
img = Image.open("my_image.jpg")  # 指向你自己的图像
predict(zero_shot_model, img, ["狗", "猫", "房子"])

# 示例输出:
# 标签概率:  {'狗': tensor([0.80590]), '猫': tensor([0.0971]), '房子': tensor([0.0970])}
# 预测标签: 狗

</details> <details> <summary>MAE训练示例</summary>

import torch
from torch.utils.data import DataLoader
from torchmultimodal.models.masked_auto_encoder.model import vit_l_16_image_mae
from torchmultimodal.models.masked_auto_encoder.utils import (
  CosineWithWarmupAndLRScaling,
)
from torchmultimodal.modules.losses.reconstruction_loss import ReconstructionLoss
from torchmultimodal.transforms.mae_transform import ImagePretrainTransform

mae_transform = ImagePretrainTransform()
dataset = MyDatasetClass(transforms=mae_transform)  # 你应该定义这个
dataloader = DataLoader(dataset, batch_size=8)

# 实例化模型和损失
mae_model = vit_l_16_image_mae()
mae_loss = ReconstructionLoss()

# 定义优化器和学习率调度器
optimizer = torch.optim.AdamW(mae_model.parameters())
lr_scheduler = CosineWithWarmupAndLRScaling(
  optimizer, max_iters=1000, warmup_iters=100  # 你应该设置这些
)

# 训练一个epoch
for batch in dataloader:
  model_out = mae_model(batch["images"])
  loss = mae_loss(model_out.decoder_pred, model_out.label_patches, model_out.mask)
  loss.backward()
  optimizer.step()
  lr_scheduler.step()

</details>

代码概览

torchmultimodal/diffusion_labs

diffusion_labs包含用于构建扩散模型的组件。有关这些组件的更多详细信息，请参阅diffusion_labs/README.md。

torchmultimodal/models

在这里可以找到模型类以及特定于给定架构的任何其他建模代码。例如，torchmultimodal/models/blip2目录包含特定于BLIP-2的建模组件。

torchmultimodal/modules

在这里可以找到可以组合在一起构建新架构的通用基本构建块。这包括层，如码本、图像块嵌入或transformer编码器/解码器，损失函数，如带温度的对比损失或重构损失，编码器，如ViT和BERT，以及融合模块，如Deep Set融合。

torchmultimodal/transforms

在这里可以找到流行模型的常用数据转换，例如CLIP、FLAVA和MAE。

安装

TorchMultimodal需要Python >= 3.8。该库可以安装带有或不带有CUDA支持的版本。 以下假设已安装conda。

先决条件

1. 安装conda环境

   conda create -n torch-multimodal python=<python_版本>
   conda activate torch-multimodal
   

2. 安装pytorch、torchvision和torchaudio。参见PyTorch文档。

   # 使用当前CUDA版本，如[此处](https://pytorch.org/get-started/locally/)所示
   # 选择每日构建的Pytorch版本，Linux作为操作系统，以及conda。选择最新的CUDA版本。
   conda install pytorch torchvision torchaudio pytorch-cuda=<cuda_版本> -c pytorch-nightly -c nvidia
   
   # 仅CPU安装
   conda install pytorch torchvision torchaudio cpuonly -c pytorch-nightly
   

从二进制文件安装

Linux上适用于Python 3.8和3.9的每日二进制文件可以通过pip wheels安装。 目前我们仅通过PyPI支持Linux平台。

python -m pip install torchmultimodal-nightly

从源代码构建

或者，您也可以从我们的源代码构建并运行我们的示例：

git clone --recursive https://github.com/facebookresearch/multimodal.git multimodal
cd multimodal

pip install -e .

对于开发者，请遵循开发安装指南。

贡献

我们欢迎来自社区的任何功能请求、错误报告或拉取请求。有关如何提供帮助，请参阅CONTRIBUTING文件。

许可证

TorchMultimodal采用BSD许可证，详见LICENSE文件。