重要

您正在查看 NeMo 2.0 文档。此版本引入了对 API 的重大更改和一个新的库，NeMo Run。我们目前正在将所有功能从 NeMo 1.0 移植到 2.0。有关先前版本或 2.0 中尚未提供的功能的文档，请参考 NeMo 24.07 文档。

嵌入器

基类

class nemo_curator.image.embedders.ImageEmbedder(

model_name: str,

image_embedding_column: str,

classifiers: Iterable[ImageClassifier],

)

用于生成图像嵌入的抽象基类。

子类只需要定义如何加载模型以及如何从 tar 文件分片中读取数据集。此类处理跨 worker 分配任务并将元数据保存到数据集。嵌入模型必须能够适应单个 GPU。

abstract load_dataset_shard(tar_path: str) → Iterable

从数据集中的 tar 文件加载图像和元数据。

参数:

tar_path (str) – 输入 WebDataset 中 tar 文件分片的路径。

返回:

数据集的迭代器。每次迭代应生成

(图像，元数据) 对的元组。图像批次将直接传递给由 ImageEmbedder.load_embedding_model 创建的模型。元数据必须是字典列表。列表的每个元素必须对应于批次中相同位置的图像。每个字典必须包含一个与数据集中 id_field 相同的字段。元数据中的此 ID 字段将用于将图像与其在元数据 (Parquet) 文件中的记录匹配。

返回类型:

Iterable

abstract load_embedding_model(device: str) → Callable

加载用于生成图像嵌入的模型。

参数:

device (str) – 一个 PyTorch 设备标识符，用于指定在哪个 GPU 上加载模型。

返回:

一个可调用模型，通常是 torch.nn.Module。

此模型的输入将是由 ImageEmbedder.load_dataset_shard 输出的图像批次。

返回类型:

Callable

Timm

class nemo_curator.image.embedders.TimmImageEmbedder(

model_name: str,

pretrained: bool = False,

batch_size: int = 1,

num_threads_per_worker: int = 4,

image_embedding_column: str = 'image_embedding',

normalize_embeddings: bool = True,

classifiers: Iterable = [],

autocast: bool = True,

use_index_files: bool = False,

)

PyTorch 图像模型 (timm) 是一个包含 SOTA 计算机视觉模型的库。其中许多模型在为 NeMo Curator 中的模块生成图像嵌入时非常有用。此模块还可以自动将图像变换从 PyTorch 变换转换为受支持模型中的 DALI 变换。

load_dataset_shard(tar_path: str)

使用 DALI 加载 WebDataset tar 分片。

参数:

tar_path (str) – 要加载的 tar 分片的路径。

返回:

数据集的迭代器。每个 tar 文件每个记录必须有 3 个文件：一个 .jpg 文件、一个 .txt 文件和一个 .json 文件。.jpg 文件必须包含图像，.txt 文件必须包含关联的标题，.json 必须包含记录的元数据（包括其 ID）。图像将使用 DALI 加载。

返回类型:

Iterable

load_embedding_model(device='cuda')

加载用于生成图像嵌入的模型。

参数:

device (str) – 一个 PyTorch 设备标识符，用于指定在哪个 GPU 上加载模型。

返回:

在指定设备上加载的 timm 模型。

如果构造函数中指定了，模型的正向调用可以使用 torch.autocast() 或嵌入归一化进行增强。

返回类型:

Callable