Hugging Face 徽标

扩散文档

UNet2DModel

Hugging Face's logo
加入 Hugging Face 社区

体验增强文档

开始

UNet 2D 模型

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

该论文的摘要如下：

众所周知，对深度神经网络进行成功的训练需要数千个标注文本样本。在本文中，我们将介绍一个网络和训练策略，依靠大幅使用数据扩充功能来更高效地利用可用的标注样本。此架构由用于获取环境的收缩路径和用于实现精确本地化的对称扩展路径组成。我们证明了即使只有数量极少的图像，此类网络也能从头到尾接受训练，并且在用于分割电子显微镜堆叠中的神经元结构的 ISBI 挑战中超越了之前的最佳方法（滑动窗口卷积神经网络）。使用同样在透射光显微镜图像（相衬和 DIC）上接受训练的网络，我们在这些类别中以较大优势赢得了 2015 年的 ISBI 细胞追踪挑战。此外，此网络运算速度非常快。在最新款 GPU 上，分割一张 512x512 像素的图像用时不到一秒。完整实现（基于 Caffe）和已接受训练的网络可在 http://lmb.informatik.uni-freiburg.de/people/ronneber/u-net 获取。

UNet2DModel

class diffusers.UNet2DModel

< >

( sample_size: Union = None in_channels: int = 3 out_channels: int = 3 center_input_sample: bool = False time_embedding_type: str = 'positional' freq_shift: int = 0 flip_sin_to_cos: bool = True down_block_types: Tuple = ('DownBlock2D', 'AttnDownBlock2D', 'AttnDownBlock2D', 'AttnDownBlock2D') up_block_types: Tuple = ('AttnUpBlock2D', 'AttnUpBlock2D', 'AttnUpBlock2D', 'UpBlock2D') block_out_channels: Tuple = (224, 448, 672, 896) layers_per_block: int = 2 mid_block_scale_factor: float = 1 downsample_padding: int = 1 downsample_type: str = 'conv' upsample_type: str = 'conv' dropout: float = 0.0 act_fn: str = 'silu' attention_head_dim: Optional = 8 norm_num_groups: int = 32 attn_norm_num_groups: Optional = None norm_eps: float = 1e-05 resnet_time_scale_shift: str = 'default' add_attention: bool = True class_embed_type: Optional = None num_class_embeds: Optional = None num_train_timesteps: Optional = None )

参数

  • sample_size (int 或 Tuple[int, int]，可选，默认为 None) — 输入/输出样本的高度和宽度。维度必须是 2 ** (len(block_out_channels) - 1) 的倍数。
  • 输出通道 (int, 可选，默认为 3) — 输出中的通道数。
  • 居中输入样本 (bool, 可选，默认为 False) — 是否将输入样本居中。
  • freq_shift (int, 可选, 默认为 0) — 傅立叶时间嵌入的频移。
  • flip_sin_to_cos (bool, 可选, 默认为 True) — 是否为傅立叶时间嵌入翻转正弦到余弦。
  • mid_block_type (str, 可选, 默认为 "UNetMidBlock2D") — UNet 中间的块类型，可以是 UNetMidBlock2D 或 UnCLIPUNetMidBlock2D。
  • up_block_types (tuple[str], 可选, 默认为 ("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D")) — 上采样块类型的元组。
  • block_out_channels (Tuple[int]，可选，默认为 (224, 448, 672, 896)) — 块输出通道的元组。
  • layers_per_block (int，可选，默认为 2) — 每个块的层数。
  • downsample_padding (int, 可选，默认为 1) — 下采样卷积的填充。
  • downsample_type (str, 可选，默认为 conv) — 下采样层的下采样类型。在 “conv” 和 “resnet” 之间选择
  • dropout (float, 可选，默认为 0.0) — 要使用的 dropout 概率。
  • act_fn (str, 可选，默认为 "silu") — 要使用的激活函数。
  • norm_num_groups (int, 可选，默认为 32) — 归一化分组数。
  • attn_norm_num_groups (int, 可选，默认为 None) — 如果设置为整数，则在中间块的 Attention 层中创建一个组归一化层，其中具有给定的组数。如果保留为 None，则仅当 resnet_time_scale_shift 设置为 default 时才会创建组归一化层，且如果创建，则具有 norm_num_groups 个组。
  • resnet_time_scale_shift (str，可选，默认为 "default") — ResNet 块的时间尺度偏移配置（请参阅 ResnetBlock2D）。从 default 或 scale_shift 中选择。
  • class_embed_type (str，可选，默认为 None) — 要使用的类嵌入类型，它最终会与时间嵌入求和。从 None、"timestep" 或 "identity" 中选择。

  • 此模型从 ModelMixin 继承。在超类文档中查看为所有模型（如下载或保存）实现的通用方法。

    forward

    < >

    ( sample: Tensor timestep: Union class_labels: Optional = None return_dict: bool = True ) → UNet2DOutput 或 元组

    参数

    • timestep (torch.Tensor 或 float 或 int) — 对输入进行去噪的步数。
    • class_labels (torch.Tensor，可选，默认值：None) — 用于调理的可选类别标签。它们的嵌入将与步数嵌入相加。
    • return_dict (bool, 可选，默认值为 True) — 是否返回 UNet2DOutput 而不是普通元组。

    返回

    UNet2DOutput 或 tuple

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

    UNet2DModel 前向方法。

UNet2DOutput

class diffusers.models.unets.unet_2d.UNet2DOutput

< >

( sample: 张量 )

参数

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

UNet2DModel的输出。

< > 在 GitHub 上更新

←UNet1DModel UNet2DConditionModel