Hugging Face 徽标

Diffusers 文档

UNet3DConditionModel

Hugging Face's logo
加入 Hugging Face 社区

并访问增强的文档体验

开始使用

UNet3DConditionModel

最初，Ronneberger 等人提出了 UNet 模型用于生物医学图像分割，但它在 🤗 Diffusers 中也很常用，因为它输出的图像大小与输入相同。它是扩散系统中最重要的组件之一，因为它促进了实际的扩散过程。根据维度数量以及是否是条件模型，🤗 Diffusers 中有几种 UNet 模型的变体。这是一个 3D UNet 条件模型。

论文的摘要是

人们普遍认为，成功训练深度网络需要数千个带注释的训练样本。在本文中，我们提出了一种网络和训练策略，该策略强烈依赖于数据增强，以更有效地利用可用的带注释样本。该架构由一个捕获上下文的收缩路径和一个能够进行精确定位的对称扩展路径组成。我们表明，这种网络可以从极少的图像中端到端训练，并且在 ISBI 挑战赛中，在电子显微镜堆栈中神经结构分割方面优于先前最佳方法（滑动窗口卷积网络）。使用在透射光显微镜图像（相差和 DIC）上训练的相同网络，我们在这些类别中以较大优势赢得了 2015 年 ISBI 细胞跟踪挑战赛。此外，该网络速度很快。在最新的 GPU 上，512x512 图像的分割不到一秒钟。完整的实现（基于 Caffe）和训练后的网络可在 http://lmb.informatik.uni-freiburg.de/people/ronneber/u-net 获取。

UNet3DConditionModel

类 diffusers.UNet3DConditionModel

< >

( sample_size: Optional = None in_channels: int = 4 out_channels: int = 4 down_block_types: Tuple = ('CrossAttnDownBlock3D', 'CrossAttnDownBlock3D', 'CrossAttnDownBlock3D', 'DownBlock3D') up_block_types: Tuple = ('UpBlock3D', 'CrossAttnUpBlock3D', 'CrossAttnUpBlock3D', 'CrossAttnUpBlock3D') block_out_channels: Tuple = (320, 640, 1280, 1280) layers_per_block: int = 2 downsample_padding: int = 1 mid_block_scale_factor: float = 1 act_fn: str = 'silu' norm_num_groups: Optional = 32 norm_eps: float = 1e-05 cross_attention_dim: int = 1024 attention_head_dim: Union = 64 num_attention_heads: Union = None time_cond_proj_dim: Optional = None )

参数

  • sample_size (int 或 Tuple[int, int], 可选，默认为 None) — 输入/输出样本的高度和宽度。
  • in_channels (int, 可选，默认为 4) — 输入样本的通道数。
  • out_channels (int, 可选，默认为 4) — 输出的通道数。
  • down_block_types (Tuple[str], 可选，默认为 ("CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "DownBlock3D")) — 要使用的下采样块的元组。
  • up_block_types (Tuple[str], 可选，默认为 ("UpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D")) — 要使用的上采样块的元组。
  • block_out_channels (Tuple[int], 可选，默认为 (320, 640, 1280, 1280)) — 每个块的输出通道的元组。
  • layers_per_block (int, 可选，默认为 2) — 每个块的层数。
  • downsample_padding (int, 可选，默认为 1) — 用于下采样卷积的填充。
  • mid_block_scale_factor (float, 可选，默认为 1.0) — 用于中间块的缩放因子。
  • norm_num_groups (int, 可选, 默认为 32) — 正则化的组数。如果为 None，则在后处理中跳过正则化和激活层。
  • norm_eps (float, 可选, 默认为 1e-5) — 正则化使用的 epsilon。
  • cross_attention_dim (int, 可选, 默认为 1024) — 交叉注意力特征的维度。
  • attention_head_dim (int, 可选, 默认为 64) — 注意力头的维度。
  • num_attention_heads (int, 可选) — 注意力头的数量。
  • time_cond_proj_dim (int, 可选, 默认为 None) — 时间步长嵌入中 cond_proj 层的维度。

一个条件 3D UNet 模型，它接收一个噪声样本、条件状态和一个时间步长，并返回一个形状输出的样本。

此模型继承自 ModelMixin。查看超类文档以了解为所有模型实现的通用方法（例如下载或保存）。

disable_freeu

< >

( )

禁用 FreeU 机制。

enable_forward_chunking

< >

( chunk_size: Optional = None dim: int = 0 )

参数

  • chunk_size (int, 可选) — 前馈层的块大小。如果未指定，则将分别在每个维度为 dim 的张量上运行前馈层。
  • dim (int, 可选, 默认为 0) — 应该对前馈计算进行分块的维度。在 dim=0（批次）或 dim=1（序列长度）之间选择。

将注意力处理器设置为使用 前馈分块。

enable_freeu

( s1 s2 b1 b2 )

参数

  • s1 (float) — 第 1 阶段的缩放因子，用于衰减跳跃特征的贡献。这旨在缓解增强去噪过程中出现的“过度平滑效应”。
  • s2 (float) — 第 2 阶段的缩放因子，用于衰减跳跃特征的贡献。这旨在缓解增强去噪过程中出现的“过度平滑效应”。
  • b1 (float) — 第 1 阶段的缩放因子，用于放大主干特征的贡献。
  • b2 (float) — 第 2 阶段的缩放因子，用于放大主干特征的贡献。

启用来自 https://arxiv.org/abs/2309.11497 的 FreeU 机制。

缩放因子后面的后缀表示它们应用到的阶段块。

请参考 官方仓库，了解适用于不同流水线（例如 Stable Diffusion v1、v2 和 Stable Diffusion XL）的有效值组合。

正向传播

< >

( sample: Tensor timestep: Union encoder_hidden_states: Tensor class_labels: Optional = None timestep_cond: Optional = None attention_mask: Optional = None cross_attention_kwargs: Optional = None down_block_additional_residuals: Optional = None mid_block_additional_residual: Optional = None return_dict: bool = True ) → UNet3DConditionOutput 或 tuple

参数

  • sample (torch.Tensor) — 具有以下形状的噪声输入张量 (batch, num_channels, num_frames, height, width。
  • timestep (torch.Tensor 或 float 或 int) — 用于对输入进行去噪的时间步数。
  • encoder_hidden_states (torch.Tensor) — 形状为 (batch, sequence_length, feature_dim) 的编码器隐藏状态。
  • class_labels (torch.Tensor, 可选, 默认值为 None) — 用于条件控制的可选类别标签。它们的嵌入将与时间步长嵌入求和。timestep_cond — (torch.Tensor, 可选, 默认值为 None): 时间步长的条件嵌入。如果提供，则这些嵌入将与通过 self.time_embedding 层传递的样本求和，以获得时间
  • attention_mask (torch.Tensor, 可选, 默认为 None) — 形状为 (batch, key_tokens) 的注意力掩码应用于 encoder_hidden_states。如果为 1，则保留掩码，否则如果为 0，则将其丢弃。掩码将转换为偏差，这会向对应于“丢弃”标记的注意力分数添加较大的负值。
  • cross_attention_kwargs (dict, 可选) — 如果指定，则传递给 diffusers.models.attention_processor 中 self.processor 下定义的 AttentionProcessor 的关键字参数字典。down_block_additional_residuals — (tuple of torch.Tensor, 可选): 如果指定，则将一组张量添加到下采样 Unet 块的残差中。mid_block_additional_residual — (torch.Tensor, 可选): 如果指定，则将一个张量添加到中间 Unet 块的残差中。
  • return_dict (bool, 可选, 默认为 True) — 是否返回 UNet3DConditionOutput 而不是普通元组。
  • cross_attention_kwargs (dict, 可选) — 如果指定，则传递给 AttnProcessor 的关键字参数字典。

返回

UNet3DConditionOutput 或 tuple

如果 return_dict 为 True，则返回 UNet3DConditionOutput，否则返回一个 tuple，其中第一个元素是样本张量。

UNet3DConditionModel 的前向方法。

fuse_qkv_projections

< >

( )

启用融合的 QKV 投影。对于自注意力模块，所有投影矩阵（即查询、键、值）都融合在一起。对于交叉注意力模块，键和值投影矩阵融合在一起。

此 API 处于 🧪 实验阶段。

set_attention_slice

< >

( slice_size: Union )

参数

  • slice_size (str 或 int 或 list(int), 可选, 默认为 "auto") — 当为 "auto" 时，注意力头的输入将减半，因此注意力将分两步计算。如果为 "max"，则通过一次仅运行一个切片来节省最大内存。如果提供了一个数字，则使用与 attention_head_dim // slice_size 相同数量的切片。在这种情况下，attention_head_dim 必须是 slice_size 的倍数。

启用切片注意力计算。

启用此选项后，注意力模块会将输入张量分割成切片，以分多个步骤计算注意力。这对于以牺牲少量速度为代价来节省一些内存很有用。

set_attn_processor

< >

( processor: Union )

参数

  • processor (dict of AttentionProcessor 或仅 AttentionProcessor) — 将被设置为**所有**Attention层的处理器的实例化处理器类或处理器类字典。

    如果 processor 是一个字典，则键需要定义对应交叉注意力处理器的路径。在设置可训练注意力处理器时，强烈建议使用此方法。

设置用于计算注意力的注意力处理器。

set_default_attn_processor

< >

( )

禁用自定义注意力处理器并设置默认注意力实现。

unfuse_qkv_projections

< >

( )

如果启用了融合的 QKV 投影，则禁用它。

此 API 处于 🧪 实验阶段。

UNet3DConditionOutput

类 diffusers.models.unets.unet_3d_condition.UNet3DConditionOutput

< >

( sample: 张量 )

参数

  • sample (torch.Tensor 形状为 (batch_size, num_channels, num_frames, height, width)) — 基于 encoder_hidden_states 输入的条件输出的隐藏状态。模型最后一层的输出。

UNet3DConditionModel 的输出。

< > 在 GitHub 上更新

←UNet2DConditionModel UNetMotionModel→