Hugging Face's logo

Diffusers documentation

UNet2DModel

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.43.0).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

UNet2DModel

Модель UNet была первоначально представлена Ronneberger et al. для сегментации биомедицинских изображений, однако она также широко используется в 🤗 Diffusers, поскольку на выходе она генерирует изображения того же размера, что и входные. Это один из ключевых компонентов диффузионной системы, так как именно он обеспечивает непосредственный процесс диффузии. В 🤗 Diffusers существует несколько вариантов модели UNet, которые различаются количеством измерений и тем, является ли модель условной. Здесь представлена 2D-модель UNet.

Аннотация из статьи:

Существует широкое согласие в том, что успешное обучение глубоких сетей требует тысяч размеченных обучающих примеров. В данной статье мы представляем сеть и стратегию обучения, которые опираются на активное использование аугментации данных для более эффективного использования доступных размеченных примеров. Архитектура состоит из сжимающего пути для захвата контекста и симметричного расширяющего пути, обеспечивающего точную локализацию. Мы показываем, что такая сеть может быть обучена end-to-end по очень малому количеству изображений и превосходит предыдущий лучший метод (сверточная сеть со скользящим окном) на конкурсе ISBI по сегментации нейронных структур в электронных микроскопических стопках. Используя ту же сеть, обученную на изображениях, полученные с помощью световой микроскопии (фазовый контраст и DIC), мы с большим отрывом выиграли конкурс ISBI по отслеживанию клеток 2015 года в этих категориях. Кроме того, сеть работает быстро. Сегментация изображения 512x512 занимает менее секунды на современном GPU. Полная реализация (на основе Caffe) и обученные сети доступны по адресу http://lmb.informatik.uni-freiburg.de/people/ronneber/u-net.

UNet2DModel

класс diffusers.UNet2DModel

< >

( 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) — Отступ для свёртки downsampling.
  • downsample_type (str, необязательный, по умолчанию conv) — Тип downsampling для слоёв downsampling. Выберите между «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, слой group norm будет создан только в том случае, если resnet_time_scale_shift установлено в значение default, и при создании будет содержать norm_num_groups групп.
  • norm_eps (float, необязательный, по умолчанию 1e-5) — Эпсилон для нормализации.
  • resnet_time_scale_shift (str, необязательный, по умолчанию "default") — конфигурация сдвига масштаба времени для блоков ResNet (см. ResnetBlock2D). Выберите default или scale_shift.
  • class_embed_type (str, необязательный, по умолчанию None) — Тип классовой эмбеддинга, который в итоге суммируется с временными эмбеддингами. Выберите одно из значений: None, "timestep" или "identity".
  • num_class_embeds (int, необязательный, по умолчанию None) — Входное измерение обучаемой матрицы эмбеддингов, которая проецируется в time_embed_dim при выполнении условного кодирования по классу, когда class_embed_type равно None.

Двумерная модель UNet, которая принимает зашумленный образец и шаг времени и возвращает вывод той же формы.

Эта модель наследуется от ModelMixin. См. документацию суперкласса для его универсальных методов, реализованных для всех моделей (например, загрузка или сохранение).

forward

< >

( 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, где первый элемент — это тензор выборки.

Метод forward модели UNet2DModel.

UNet2DOutput

класс diffusers.models.unets.unet_2d.UNet2DOutput

< >

( sample: Tensor )

Параметры

  • sample (torch.Tensor формы (batch_size, num_channels, height, width)) — Скрытые состояния, выведенные последним слоем модели.

Вывод UNet2DModel.

Update on GitHub

←UNet2DConditionModel UNet3DConditionModel→