重要提示

您正在查看 NeMo 2.0 文档。此版本对 API 和新库 NeMo Run 进行了重大更改。我们目前正在将 NeMo 1.0 的所有功能移植到 2.0。有关先前版本或 2.0 中尚不可用的功能的文档，请参阅 NeMo 24.07 文档。

NeMo 多模态 API#

模型类#

class nemo.collections.nlp.models.language_modeling.megatron_base_model.MegatronBaseModel(*args: Any, **kwargs: Any)

基类： NLPModel

Megatron 基类。所有 NeMo Megatron 模型都从此类继承。

  • 为 nemo 启用模型并行世界。

  • 开启所有 NVIDIA 优化。

  • 如果 cfg.tokenizer 可用，它将加载分词器并将词汇表填充到张量模型并行的正确大小。

  • 如果使用分布式优化器，请配置为与 O2 级别优化和/或模型并行性兼容。

  • 执行梯度裁剪：grad_clip_pl_default 触发 PyTorch Lightning 默认实现，with_distributed_adam 触发分布式优化器的实现，megatron_amp_O2 触发主梯度上的梯度裁剪，否则在模型梯度上执行梯度裁剪。

__init__(
cfg: omegaconf.dictconfig.DictConfig,
trainer: lightning.pytorch.trainer.trainer.Trainer,
no_lm_init=True,
)

所有 NeMo 模型应继承的基类

参数:

  • cfg (DictConfig) –

    配置对象。cfg 对象应具有（可选）以下子配置

    • train_ds - 用于实例化训练数据集

    • validation_ds - 用于实例化验证数据集

    • test_ds - 用于实例化测试数据集

    • optim - 用于实例化带有学习率调度器的优化器

  • trainer (Optional) – Pytorch Lightning Trainer 实例

class nemo.collections.multimodal.models.text_to_image.stable_diffusion.ldm.ddpm.MegatronLatentDiffusion(*args: Any, **kwargs: Any)#

基类： NLPAdapterModelMixin, MegatronBaseModel

Megatron 潜在扩散模型。

__init__(
cfg: omegaconf.DictConfig,
trainer: lightning.pytorch.Trainer,
)#

所有 NeMo 模型应继承的基类

参数:

  • cfg (DictConfig) –

    配置对象。cfg 对象应具有（可选）以下子配置

    • train_ds - 用于实例化训练数据集

    • validation_ds - 用于实例化验证数据集

    • test_ds - 用于实例化测试数据集

    • optim - 用于实例化带有学习率调度器的优化器

  • trainer (Optional) – Pytorch Lightning Trainer 实例

setup(stage=None)#

在 DDP 生成后执行的 PTL 钩子。

我们在此处设置数据集，因为 Megatron 数据集需要 DDP 来实例化。有关更多信息，请参阅 PyTorch Lightning 文档： https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#setup

参数:

stage (str, optional) – 可以是 ‘fit’、‘validate’、‘test’ 或 ‘predict’。默认为 None。

training_step(batch)#

注意：training_step 过去具有以下签名以支持流水线并行

def training_step(self, dataloader_iter, batch_idx)

但是，完整的迭代 CUDA Graph 回调与此签名不兼容，因为我们需要包装数据加载器以在 CUDA Graph 外部生成静态张量。此签名将 next(dataloader) 移动到 CUDA Graph 捕获区域，因此我们禁用了它。

我们的数据加载器生成微批次，然后我们从数据加载器中获取多个微批次，具体取决于全局批次大小和模型并行大小，以生成微批次列表。批次应为微批次列表，并且这些微批次应在 CPU 上。然后，在流水线期间将微批次移动到 GPU。然后，使用 Apex fwd/bwd 函数通过流水线管道传输微批次列表。

class nemo.collections.multimodal.models.text_to_image.dreambooth.dreambooth.MegatronDreamBooth(*args: Any, **kwargs: Any)#

基类： NLPAdapterModelMixin, MegatronBaseModel

__init__(
cfg: omegaconf.DictConfig,
trainer: lightning.pytorch.Trainer,
)#

所有 NeMo 模型应继承的基类

参数:

  • cfg (DictConfig) –

    配置对象。cfg 对象应具有（可选）以下子配置

    • train_ds - 用于实例化训练数据集

    • validation_ds - 用于实例化验证数据集

    • test_ds - 用于实例化测试数据集

    • optim - 用于实例化带有学习率调度器的优化器

  • trainer (Optional) – Pytorch Lightning Trainer 实例

setup(stage=None)#

在 DDP 生成后执行的 PTL 钩子。

我们在此处设置数据集，因为 Megatron 数据集需要 DDP 来实例化。有关更多信息，请参阅 PyTorch Lightning 文档： https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#setup

参数:

stage (str, optional) – 可以是 ‘fit’、‘validate’、‘test’ 或 ‘predict’。默认为 None。

training_step(dataloader_iter)#

我们的数据加载器生成微批次，然后我们从数据加载器中获取多个微批次，具体取决于全局批次大小和模型并行大小，以生成微批次列表。批次应为微批次列表，并且这些微批次应在 CPU 上。然后，在流水线期间将微批次移动到 GPU。然后，使用 Apex fwd/bwd 函数通过流水线管道传输微批次列表。

class nemo.collections.multimodal.models.text_to_image.controlnet.controlnet.MegatronControlNet(*args: Any, **kwargs: Any)#

基类： MegatronBaseModel

__init__(
cfg: omegaconf.DictConfig,
trainer: lightning.pytorch.Trainer,
)#

所有 NeMo 模型应继承的基类

参数:

  • cfg (DictConfig) –

    配置对象。cfg 对象应具有（可选）以下子配置

    • train_ds - 用于实例化训练数据集

    • validation_ds - 用于实例化验证数据集

    • test_ds - 用于实例化测试数据集

    • optim - 用于实例化带有学习率调度器的优化器

  • trainer (Optional) – Pytorch Lightning Trainer 实例

setup(stage=None)#

在 DDP 生成后执行的 PTL 钩子。

我们在此处设置数据集，因为 Megatron 数据集需要 DDP 来实例化。有关更多信息，请参阅 PyTorch Lightning 文档： https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#setup

参数:

stage (str, optional) – 可以是 ‘fit’、‘validate’、‘test’ 或 ‘predict’。默认为 None。

training_step(dataloader_iter)#

我们的数据加载器生成微批次，然后我们从数据加载器中获取多个微批次，具体取决于全局批次大小和模型并行大小，以生成微批次列表。批次应为微批次列表，并且这些微批次应在 CPU 上。然后，在流水线期间将微批次移动到 GPU。然后，使用 Apex fwd/bwd 函数通过流水线管道传输微批次列表。

class nemo.collections.multimodal.models.text_to_image.imagen.imagen.MegatronImagen(*args: Any, **kwargs: Any)#

基类： MegatronBaseModel

__init__(
cfg: omegaconf.DictConfig,
trainer: lightning.pytorch.Trainer,
)#

所有 NeMo 模型应继承的基类

参数:

  • cfg (DictConfig) –

    配置对象。cfg 对象应具有（可选）以下子配置

    • train_ds - 用于实例化训练数据集

    • validation_ds - 用于实例化验证数据集

    • test_ds - 用于实例化测试数据集

    • optim - 用于实例化带有学习率调度器的优化器

  • trainer (Optional) – Pytorch Lightning Trainer 实例

setup(stage=None)#

在 DDP 生成后执行的 PTL 钩子。

我们在此处设置数据集，因为 Megatron 数据集需要 DDP 来实例化。有关更多信息，请参阅 PyTorch Lightning 文档： https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#setup

参数:

stage (str, optional) – 可以是 ‘fit’、‘validate’、‘test’ 或 ‘predict’。默认为 None。

training_step(dataloader_iter)#

我们的数据加载器生成微批次，然后我们从数据加载器中获取多个微批次，具体取决于全局批次大小和模型并行大小，以生成微批次列表。批次应为微批次列表，并且这些微批次应在 CPU 上。然后，在流水线期间将微批次移动到 GPU。然后，使用 Apex fwd/bwd 函数通过流水线管道传输微批次列表。

validation_step(dataloader_iter)#

我们的数据加载器生成微批次，然后我们从数据加载器中获取多个微批次，具体取决于全局批次大小和模型并行大小，以生成微批次列表。然后，使用 megatron-core fwd/bwd 函数通过流水线管道传输微批次列表。

模块#

class nemo.collections.multimodal.modules.stable_diffusion.diffusionmodules.openaimodel.UNetModel(*args: Any, **kwargs: Any)#

基类： Module

具有注意力和时间步长嵌入的完整 UNet 模型。

参数:

  • in_channels (int) – 输入张量中的通道数。

  • model_channels (int) – 模型的基本通道计数。

  • out_channels (int) – 输出张量中的通道数。

  • num_res_blocks (int) – 每个下采样的残差块数。

  • attention_resolutions (set/list/tuple) – 应用注意力的下采样率。例如，如果这包括 4，则在 4 倍下采样时使用注意力。

  • dropout (float) – dropout 概率。

  • channel_mult (list/tuple) – UNet 每个级别的通道乘数。

  • conv_resample (bool) – 如果为 True，则使用学习的卷积进行上采样和下采样。

  • dims (int) – 确定信号是 1D、2D 还是 3D。

  • num_classes (int, optional) – 如果指定，模型将变为类条件模型，并具有给定的类数。

  • use_checkpoint (bool) – 如果为 True，则使用梯度检查点以减少内存使用量。

  • num_heads (int) – 每个注意力层中的注意力头数。

  • num_heads_channels (int, optional) – 如果指定，则覆盖 num_heads 并为每个注意力头使用固定的通道宽度。

  • num_heads_upsample (int, optional) – 为上采样设置不同的头数。已弃用。

  • use_scale_shift_norm (bool) – 如果为 True，则使用类似 FiLM 的条件机制。

  • resblock_updown (bool) – 如果为 True，则对上/下采样使用残差块。

  • use_new_attention_order (bool) – 如果为 True，则使用不同的注意力模式以提高效率。

class nemo.collections.multimodal.modules.imagen.diffusionmodules.nets.UNetModel(*args: Any, **kwargs: Any)#

基类： Module

用于 Imagen Base 和 SR 模型的具有注意力和时间步长嵌入的完整 UNet 模型。

参数:

  • embed_dim – 嵌入的维度。也用于计算 ResBlock 中的通道数。

  • image_size – 输入图像大小。用于计算在 UNet 中注入注意力层的位置。

  • channels – 输入通道数，默认为 3。

  • text_embed_dim – 条件文本嵌入的维度。不同的文本编码器和不同的模型版本具有不同的值，默认为 512

  • num_res_blocks – UNet 每个级别的 ResBlock 数量，默认为 3。

  • channel_mult – 与 embed_dim 一起使用以计算 UNet 每个级别的通道数，默认为 [1, 2, 3, 4]

  • num_attn_heads – 注意力层中的头数，默认为 4。

  • per_head_channels – 每个注意力头的通道数，默认为 64。

  • cond_dim – 条件投影的维度，默认为 512。

  • attention_type – 注意力层的类型，默认为 ‘fused’。

  • feature_pooling_type – 池化类型，默认为 ‘attention’。

  • learned_sinu_pos_emb_dim – 学习的时间位置嵌入的维度。0 表示未学习的时间步长嵌入。默认为 16

  • attention_resolutions – 注入注意力层的分辨率列表。默认为 [8, 16, 32]

  • dropout – dropout 率，默认为 0。

  • use_null_token – 是否为注意力创建学习的空标记，默认为 False。

  • init_conv_kernel_size – 初始 Conv 内核大小，默认为 3。

  • gradient_checkpointing – 是否使用梯度检查点，默认为 False。

  • scale_shift_norm – 是否使用 scale shift norm，默认为 False。

  • stable_attention – 是否使用数值稳定的注意力计算，默认为 True。

  • flash_attention – 是否使用 flash attention 计算，默认为 False。

  • resblock_updown – 是否使用 ResBlock 或 Downsample/Upsample，默认为 False。

  • resample_with_conv – 当 resblock_updown=False 时，是否除了 Pooling&ConvTranspose 外还使用 conv。默认为 True。

  • low_res_cond – 是否以低分辨率输入为条件，用于 SR 模型。默认为 False。

  • noise_cond_aug – 是否使用低分辨率输入添加噪声条件增强。默认为 False。

class nemo.collections.multimodal.modules.imagen.diffusionmodules.nets.EfficientUNetModel(*args: Any, **kwargs: Any)#

基类： Module

用于 Imagen SR 模型的具有注意力和时间步长嵌入的完整高效 UNet 模型。

参数:

  • embed_dim – 嵌入的维度。也用于计算 ResBlock 中的通道数。

  • image_size – 输入图像大小。用于计算在 UNet 中注入注意力层的位置。

  • channels – 输入通道数，默认为 3。

  • text_embed_dim – 条件文本嵌入的维度。不同的文本编码器和不同的模型版本具有不同的值，默认为 512

  • channel_mult – 与 embed_dim 一起使用以计算 UNet 每个级别的通道数，默认为 [1, 1, 2, 4, 8]。

  • num_attn_heads – 注意力层中的头数，默认为 8。

  • per_head_channels – 每个注意力头的通道数，默认为 64。

  • attention_type – 注意力层的类型，默认为 ‘fused’。

  • atnn_enabled_at – 是否在每个级别启用注意力，默认为 [0, 0, 0, 0, 1]。

  • feature_pooling_type – 池化类型，默认为 ‘attention’。

  • stride – ResBlock 中的步幅，默认为 2。

  • num_resblocks – 与 num_res_blocks 一起使用以计算高效 UNet 每个级别的残差块数。默认为 [1, 2, 4, 8, 8]。

  • learned_sinu_pos_emb_dim – 学习的时间位置嵌入的维度。0 表示未学习的时间步长嵌入。默认为 16

  • use_null_token – 是否为注意力创建学习的空标记，默认为 False。

  • init_conv_kernel_size – 初始 Conv 内核大小，默认为 3。

  • gradient_checkpointing – 是否使用梯度检查点，默认为 False。

  • scale_shift_norm – 是否使用 scale shift norm，默认为 False。

  • stable_attention – 是否使用数值稳定的注意力计算，默认为 True。

  • flash_attention – 是否使用 flash attention 计算，默认为 False。

  • skip_connection_scaling – 是否对 ResBlock 跳跃连接使用 1/sqrt(2) 缩放，默认为 False。

  • noise_cond_aug – 是否使用低分辨率输入添加噪声条件增强。默认为 False。

class nemo.collections.multimodal.models.text_to_image.stable_diffusion.ldm.autoencoder.AutoencoderKL(*args: Any, **kwargs: Any)#

基类： LightningModule

__init__(
ddconfig,
embed_dim,
lossconfig=None,
ckpt_path=None,
ignore_keys=[],
image_key='image',
colorize_nlabels=None,
from_NeMo=False,
monitor=None,
from_pretrained: str | None = None,
)#
decode(z)#

将潜在表示解码回像素空间。

encode(x)#

将像素空间中的输入图像编码为潜在表示。

class nemo.collections.multimodal.modules.stable_diffusion.encoders.modules.FrozenMegatronCLIPEmbedder(*args: Any, **kwargs: Any)#

基类： AbstractEmbModel

__init__(
restore_from_path,
device='cuda',
layer='last',
freeze=True,
cfg=None,
always_return_pooled=False,
enable_lora_finetune=False,
)#
forward(text)#

从输入文本获取嵌入

class nemo.collections.multimodal.modules.imagen.encoder.t5encoder.T5Encoder(*args: Any, **kwargs: Any)#

基类： Module

__init__(max_seq_len=512, encoder_path=None)#

初始化 T5 编码器。

参数:

  • max_seq_len – 最大 token 长度，默认为 512

  • encoder_path – 如果在磁盘上加载 T5，则为可选，默认为 None

encode(text_batch, device='cuda')#

将一批文本编码为 T5 嵌入。

class nemo.collections.multimodal.models.text_to_image.controlnet.controlnet.ControlledUnetModel(*args: Any, **kwargs: Any)#

基类： UNetModel

修改后的 Unet 类，在前向传播过程中结合了控制副本和冻结副本的输出。

forward(
x,
timesteps=None,
context=None,
control=None,
only_mid_control=False,
**kwargs,
)#
参数:

  • x – 扩散过程的潜在变量

  • timesteps – 扩散步骤

  • context – 指导去噪过程的文本嵌入

  • control – 来自每个对应层的控制副本的输出

  • only_mid_control – 是否仅添加来自中间块的控制副本的输出

数据集#

class nemo.collections.multimodal.data.common.webdataset.WebDatasetCommon(*args: Any, **kwargs: Any)#

基类: IterableDataset

NeMo 多模态模型共享的通用数据集对象。

class nemo.collections.multimodal.data.dreambooth.dreambooth_dataset.DreamBoothDataset(*args: Any, **kwargs: Any)#

基类: Dataset

用于准备实例和类别图像以及提示以微调模型的数据集。它预处理图像并标记化提示。

参数:

  • instance_data_root – 必需，包含对象图像文件的目录

  • instance_prompt – 带有与实例图像关联的特殊 token 的 captions（描述）

  • with_prior_preservation – 是否使用来自骨干网络的原始推理输出正则化模型微调

  • reg_data_root – 用于保存来自骨干网络的推理图像的目录

  • reg_prompt – 用于生成正则化图像的 prompt（提示）

  • size – 调整图像大小以用于训练数据管道

  • center_crop – 是否对输入图像执行中心裁剪

  • load_cache_latents – 设置为 True 时，图像将转换为缓存的潜在变量，这些变量将直接加载用于训练

  • vae – vae 实例，用于将图像从像素空间编码到潜在空间

本页目录