Diffusers documentation
UNet2DModel
UNet2DModel
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
< source >( sample_size: int | tuple[int, int] | None = Nonein_channels: int = 3out_channels: int = 3center_input_sample: bool = Falsetime_embedding_type: str = 'positional'time_embedding_dim: int | None = Nonefreq_shift: int = 0flip_sin_to_cos: bool = Truedown_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 = 2mid_block_scale_factor: float = 1downsample_padding: int = 1downsample_type: str = 'conv'upsample_type: str = 'conv'dropout: float = 0.0act_fn: str = 'silu'attention_head_dim: int | None = 8norm_num_groups: int = 32attn_norm_num_groups: int | None = Nonenorm_eps: float = 1e-05resnet_time_scale_shift: str = 'default'add_attention: bool = Trueclass_embed_type: str | None = Nonenum_class_embeds: int | None = Nonenum_train_timesteps: int | None = 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]，可选，默认为("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) — 如果设置为整数，将在中间块的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。请查阅其父类文档，了解为所有模型实现的通用方法（例如下载或保存）。
forward
< source >( sample: Tensortimestep: typing.Union[torch.Tensor, float, int]class_labels: typing.Optional[torch.Tensor] = Nonereturn_dict: bool = True ) → UNet2DOutput or tuple
参数
- sample (
torch.Tensor) — 形状为(batch, channel, height, width)的含噪输入张量。 - 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
< source >( sample: Tensor )
UNet2DModel 的输出。