Hugging Face 的徽标

Diffusers 文档

UNet2D模型

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

UNet2DModel

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

该论文的摘要是

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

UNet2DModel

class diffusers.UNet2DModel

< >

( sample_size: typing.Union[int, typing.Tuple[int, int], NoneType] = None in_channels: int = 3 out_channels: int = 3 center_input_sample: bool = False time_embedding_type: str = 'positional' time_embedding_dim: typing.Optional[int] = None freq_shift: int = 0 flip_sin_to_cos: bool = True down_block_types: typing.Tuple[str, ...] = ('DownBlock2D', 'AttnDownBlock2D', 'AttnDownBlock2D', 'AttnDownBlock2D') up_block_types: typing.Tuple[str, ...] = ('AttnUpBlock2D', 'AttnUpBlock2D', 'AttnUpBlock2D', 'UpBlock2D') block_out_channels: typing.Tuple[int, ...] = (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: typing.Optional[int] = 8 norm_num_groups: int = 32 attn_norm_num_groups: typing.Optional[int] = None norm_eps: float = 1e-05 resnet_time_scale_shift: str = 'default' add_attention: bool = True class_embed_type: typing.Optional[str] = None num_class_embeds: typing.Optional[int] = None num_train_timesteps: typing.Optional[int] = None )

参数

  • sample_size (int 或 Tuple[int, int], 可选, 默认为 None) — 输入/输出样本的高度和宽度。 维度必须是 2 ** (len(block_out_channels) - 1) 的倍数。
  • 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) — 是否翻转傅里叶时间嵌入的 sin 为 cos。
  • down_block_types (Tuple[str], optional, defaults to ("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D")) — 下采样模块类型的元组，默认为 ("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D")。
  • mid_block_type (str, optional, defaults to "UNetMidBlock2D") — UNet 中间层的模块类型，可以是 UNetMidBlock2D 或 UnCLIPUNetMidBlock2D。
  • up_block_types (Tuple[str], optional, defaults to ("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D")) — 上采样模块类型的元组，默认为 ("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D")。
  • block_out_channels (Tuple[int], optional, defaults to (224, 448, 672, 896)) — 模块输出通道的元组，默认为 (224, 448, 672, 896)。
  • layers_per_block (int, optional, defaults to 2) — 每个模块的层数，默认为 2。
  • mid_block_scale_factor (float, optional, defaults to 1) — 中间模块的缩放因子，默认为 1。
  • downsample_padding (int, optional, defaults to 1) — 下采样卷积的填充大小，默认为 1。
  • downsample_type (str, optional, defaults to conv) — 下采样层的下采样类型。可选择 “conv” 和 “resnet”。
  • upsample_type (str, optional, defaults to conv) — 上采样层的上采样类型。可选择 “conv” 和 “resnet”。
  • dropout (float, optional, defaults to 0.0) — 要使用的 dropout 概率，默认为 0.0。
  • act_fn (str, optional, defaults to "silu") — 要使用的激活函数，默认为 "silu"。
  • attention_head_dim (int, optional, defaults to 8) — 注意力头部的维度，默认为 8。
  • norm_num_groups (int, optional, defaults to 32) — 归一化的组数，默认为 32。
  • attn_norm_num_groups (int, optional, defaults to None) — 如果设置为整数，则将在中间模块的 Attention 层中创建一个组归一化层，其中包含给定数量的组。如果保留为 None，则仅当 resnet_time_scale_shift 设置为 default 时才创建组归一化层，如果创建，则将具有 norm_num_groups 组。
  • norm_eps (float, optional, defaults to 1e-5) — 归一化的 epsilon 值，默认为 1e-5。
  • resnet_time_scale_shift (str, optional, defaults to "default") — ResNet 块的时间尺度偏移配置 (参见 ResnetBlock2D)。从 default 或 scale_shift 中选择。
  • class_embed_type (str, optional, defaults to None) — 要使用的类别嵌入类型，最终与时间步嵌入求和。从 None、"timestep" 或 "identity" 中选择。
  • num_class_embeds (int, optional, defaults to None) — 当 class_embed_type 等于 None 时，可学习嵌入矩阵的输入维度，将被投影到 time_embed_dim。

一个 2D UNet 模型，它接收一个噪声样本和一个时间步，并返回一个样本形状的输出。

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

forward

< >

( sample: Tensor timestep: typing.Union[torch.Tensor, float, int] class_labels: typing.Optional[torch.Tensor] = None return_dict: bool = True ) → UNet2DOutput 或 tuple

参数

  • sample (torch.Tensor) — 噪声输入张量，形状为 (batch, channel, height, width)。
  • timestep (torch.Tensor 或 float 或 int) — 用于去噪输入的 timestep 数量。
  • class_labels (torch.Tensor, optional, defaults to None) — 用于条件控制的可选类别标签。它们的嵌入将与 timestep 嵌入求和。
  • return_dict (bool, optional, defaults to True) — 是否返回 UNet2DOutput 而不是普通元组，默认为 True。

返回值

UNet2DOutput 或 tuple

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

UNet2DModel 的 forward 方法。

UNet2DOutput

class diffusers.models.unets.unet_2d.UNet2DOutput

< >

( sample: Tensor )

参数

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

UNet2DModel 的输出。

< > 在 GitHub 上更新

←UNet1DModel UNet2DConditionModel→