)

TorchMultimodal（测试版）

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

简介

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

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

模型

TorchMultimodal 包含多种模型，包括：

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

示例脚本

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

快速入门

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

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])}

# 预测标签：狗

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()

代码概览

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 的版本或不支持 CUDA 的版本。 以下内容假设已安装 conda。

先决条件

1. 创建 conda 环境

   conda create -n torch-multimodal python=\<python_version\>
   conda activate torch-multimodal
   

2. 安装 PyTorch、Torchvision 和 Torchaudio。请参阅 PyTorch 文档。

   # 使用当前的 CUDA 版本，详见 [这里](https://pytorch.org/get-started/locally/)
   # 选择 nightly 版本的 PyTorch，操作系统为 Linux，并使用 conda。选择最新的 CUDA 版本。
   conda install pytorch torchvision torchaudio pytorch-cuda=\<cuda_version\> -c pytorch-nightly -c nvidia
   
   # 如果仅需 CPU 版本
   conda install pytorch torchvision torchaudio cpuonly -c pytorch-nightly
   

从二进制文件安装

适用于 Python 3.8 和 3.9 的 Linux 平台 nightly 二进制文件可以通过 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 文件。

TorchMultimodal 快速上手指南

TorchMultimodal 是一个基于 PyTorch 的库，旨在大规模训练最先进的多模态多任务模型（包括内容理解和生成模型）。它提供了模块化的构建块、预训练模型类以及复现前沿研究的示例代码。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux (推荐，支持夜间版二进制安装)
• Python 版本: >= 3.8
• 包管理器: Conda (推荐)
• 硬件: 支持 CUDA 的 GPU (可选，也支持 CPU 模式)

前置依赖: 您需要先安装 PyTorch 生态的基础组件（pytorch, torchvision, torchaudio）。

安装步骤

1. 创建并激活 Conda 环境

conda create -n torch-multimodal python=3.9
conda activate torch-multimodal

2. 安装 PyTorch 基础依赖

根据您的硬件环境选择以下一种方式安装。建议访问 PyTorch 官网 获取最新的 CUDA 版本对应命令。

选项 A: GPU 版本 (推荐) 请将 <cuda_version> 替换为您实际的 CUDA 版本（如 11.8 或 12.1）：

conda install pytorch torchvision torchaudio pytorch-cuda=<cuda_version> -c pytorch-nightly -c nvidia

选项 B: CPU 仅用版本

conda install pytorch torchvision torchaudio cpuonly -c pytorch-nightly

3. 安装 TorchMultimodal

目前主要通过 pip 安装夜间版（nightly）二进制包（支持 Linux Python 3.8/3.9）：

pip install torchmultimodal-nightly

注意: 如果您在中国大陆地区遇到下载速度慢的问题，建议在 pip 命令后添加国内镜像源参数，例如： pip install torchmultimodal-nightly -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

以下展示两个最核心的使用场景：使用预训练模型进行零样本推理（Zero-shot）以及构建一个简单的训练循环。

场景一：FLAVA 模型零样本图像分类

此示例演示如何加载预训练的 FLAVA 模型，对图像进行零样本分类。

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(
            "Label probabilities: ",
            {labels[i]: probs[:, i] for i in range(len(labels))},
        )
        print(f"Predicted label: {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, ["dog", "cat", "house"])

# 示例输出:
# Label probabilities:  {'dog': tensor([0.80590]), 'cat': tensor([0.0971]), 'house': tensor([0.0970])}
# Predicted label: dog

场景二：MAE 模型训练循环示例

此示例演示如何实例化 MAE 模型、损失函数及优化器，并执行一个基本的训练步。

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  # 请根据实际需求设置参数
)

# 执行单轮训练
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()