# UNet2DModel

[UNet](https://huggingface.co/papers/1505.04597) 模型最初由 Ronneberger 等人提出，用于生物医学图像分割，但由于其输出的图像尺寸与输入相同，因此在 🤗 Diffusers 中也被广泛使用。它是扩散系统中最关键的组件之一，因为它实现了实际的扩散过程。🤗 Diffusers 中存在多种 UNet 模型变体，具体取决于其维度数量以及是否为条件模型。这里介绍的是一个 2D UNet 模型。

该论文的摘要如下：

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

## UNet2DModel[[diffusers.UNet2DModel]]

#### diffusers.UNet2DModel[[diffusers.UNet2DModel]]

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

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/unets/unet_2d.py#L39)

**参数：**

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]`，*可选*，默认为 `("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D")`)：下采样块类型的元组。

mid_block_type (`str`，*可选*，默认为 `"UNetMidBlock2D"`)：UNet 中间部分的块类型，可以是 `UNetMidBlock2D` 或 `None`。

up_block_types (`tuple[str]`，*可选*，默认为 `("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D")`)：上采样块类型的元组。

block_out_channels (`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`)：下采样层的下采样类型。在 "conv" 和 "resnet" 之间选择

upsample_type (`str`，*可选*，默认为 `conv`)：上采样层使用的上采样类型。在 "conv" 和 "resnet" 之间选择

dropout (`float`，*可选*，默认值为 0.0)：要使用的 dropout 概率。

act_fn (`str`，*可选*，默认为 `"silu"`)：要使用的激活函数。

attention_head_dim (`int`，*可选*，默认为 `8`)：注意力头维度。

norm_num_groups (`int`，*可选*，默认为 `32`)：用于归一化的组数。

attn_norm_num_groups (`int`，*可选*，默认为 `None`)：如果设置为整数，将在 mid block 的 `Attention` 层中创建一个具有指定组数的 group norm 层。如果保持为 `None`，则仅当 `resnet_time_scale_shift` 设置为 `default` 时才会创建 group norm 层，且创建时其组数为 `norm_num_groups`。

norm_eps (`float`，*可选*，默认为 `1e-5`)：归一化的 epsilon 值。

resnet_time_scale_shift (`str`，*可选*，默认为 `"default"`)：ResNet 块的时间尺度偏移配置（参见 `ResnetBlock2D`）。可选值为 `default` 或 `scale_shift`。

class_embed_type (`str`，*可选*，默认为 `None`)：要使用的类嵌入类型，最终会与时间嵌入相加。可选值为 `None`、`"timestep"` 或 `"identity"`。

num_class_embeds (`int`，*可选*，默认为 `None`)：当使用 `class_embed_type` 等于 `None` 的类别条件时，可学习嵌入矩阵的输入维度，该矩阵将被投影到 `time_embed_dim`。

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

该模型继承自 [ModelMixin](/docs/diffusers/v0.43.0/zh/api/models/overview#diffusers.ModelMixin)。请查阅其父类文档，了解为所有模型实现的通用方法（例如下载或保存）。

#### forward[[diffusers.UNet2DModel.forward]]

```python
forward(sample: Tensor, timestep: typing.Union[torch.Tensor, float, int], class_labels: typing.Optional[torch.Tensor] = None, return_dict: bool = True)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/unets/unet_2d.py#L249)

**参数：**

sample (`torch.Tensor`) : 具有如下形状 `(batch, channel, height, width)` 的含噪输入张量。

timestep (`torch.Tensor` 或 `float` 或 `int`)：用于对输入进行去噪的时间步数。

class_labels (`torch.Tensor`，*可选*，默认为 `None`)：用于条件控制的可选类别标签。它们的嵌入将与时间步嵌入相加。

return_dict (`bool`，*可选*，默认为 `True`)：是否返回 [UNet2DOutput](/docs/diffusers/v0.43.0/zh/api/models/unet2d#diffusers.models.unets.unet_2d.UNet2DOutput) 而不是普通元组。

**返回：** [UNet2DOutput](/docs/diffusers/v0.43.0/zh/api/models/unet2d#diffusers.models.unets.unet_2d.UNet2DOutput) 或 `tuple`

如果 `return_dict` 为 True，则返回 [UNet2DOutput](/docs/diffusers/v0.43.0/zh/api/models/unet2d#diffusers.models.unets.unet_2d.UNet2DOutput)，否则返回一个 `tuple`，其中第一个元素是样本张量。

[UNet2DModel](/docs/diffusers/v0.43.0/zh/api/models/unet2d#diffusers.UNet2DModel) 的前向方法。

## UNet2DOutput[[diffusers.models.unets.unet_2d.UNet2DOutput]]

#### diffusers.models.unets.unet_2d.UNet2DOutput[[diffusers.models.unets.unet_2d.UNet2DOutput]]

```python
diffusers.models.unets.unet_2d.UNet2DOutput(sample: Tensor)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/unets/unet_2d.py#L27)

**参数：**

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

[UNet2DModel](/docs/diffusers/v0.43.0/zh/api/models/unet2d#diffusers.UNet2DModel) 的输出。