Diffusers 文档
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
< 来源 >( 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') mid_block_type: typing.Optional[str] = 'UNetMidBlock2D' 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], 可选, 默认为("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层中创建一个组归一化层，其组数为给定值。如果保留为None，则仅当resnet_time_scale_shift设置为default时才会创建组归一化层，如果创建，则其组数为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
< 来源 >( 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) — 去噪输入的时间步数。 - 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: Tensor )
UNet2DModel 的输出。