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

O modelo UNet foi originalmente introduzido por Ronneberger et al. para segmentação de imagens biomédicas, mas também é comumente usado no 🤗 Diffusers porque ele gera imagens com o mesmo tamanho da entrada. Ele é um dos componentes mais importantes de um sistema de difusão, pois facilita o processo de difusão em si. Existem várias variantes do modelo UNet no 🤗 Diffusers, dependendo do número de dimensões e de se trata de um modelo condicional ou não. Este é um modelo UNet 2D.

O resumo do artigo é:

Há amplo consenso de que o treinamento bem-sucedido de redes profundas requer milhares de amostras de treinamento anotadas. Neste artigo, apresentamos uma rede e uma estratégia de treinamento que dependem do uso intenso de aumento de dados para utilizar as amostras anotadas disponíveis de forma mais eficiente. A arquitetura consiste em um caminho contrativo para capturar o contexto e um caminho expansivo simétrico que permite a localização precisa. Mostramos que uma rede como essa pode ser treinada de ponta a ponta a partir de muito poucas imagens e supera o melhor método anterior (uma rede convolucional de janela deslizante) no desafio ISBI para segmentação de estruturas neuronais em pilhas de microscopia eletrônica. Usando a mesma rede treinada em imagens de microscopia de luz transmitida (contraste de fase e DIC), vencemos o desafio de rastreamento celular ISBI 2015 nessas categorias por uma grande margem. Além disso, a rede é rápida. A segmentação de uma imagem de 512x512 leva menos de um segundo em uma GPU recente. A implementação completa (baseada em Caffe) e as redes treinadas estão disponíveis em http://lmb.informatik.uni-freiburg.de/people/ronneber/u-net.

UNet2DModel

classe 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 )

Parâmetros

  • sample_size (int ou tuple[int, int], opcional, padrão None) — Altura e largura da amostra de entrada/saída. As dimensões devem ser múltiplos de 2 ** (len(block_out_channels) - 1).
  • in_channels (int, opcional, padrão 3) — Número de canais na amostra de entrada.
  • out_channels (int, opcional, padrão 3) — Número de canais na saída.
  • center_input_sample (bool, opcional, padrão False) — Se deve centralizar a amostra de entrada.
  • time_embedding_type (str, opcional, padrão "positional") — Tipo de incorporação de tempo a ser usado.
  • freq_shift (int, opcional, padrão 0) — Deslocamento de frequência para incorporação temporal de Fourier.
  • flip_sin_to_cos (bool, opcional, padrão True) — Se deve inverter sin para cos na incorporação de tempo de Fourier.
  • down_block_types (tuple[str], opcional, padrão ("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D")) — tupla de tipos de blocos de downsampling.
  • mid_block_type (str, opcional, padrão "UNetMidBlock2D") — Tipo de bloco para o meio do UNet, pode ser UNetMidBlock2D ou None.
  • up_block_types (tuple[str], opcional, padrão ("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D")) — tupla de tipos de blocos de upsample.
  • block_out_channels (tuple[int], opcional, padrão (224, 448, 672, 896)) — tupla de canais de saída dos blocos.
  • layers_per_block (int, opcional, padrão 2) — O número de camadas por bloco.
  • mid_block_scale_factor (float, opcional, padrão 1) — O fator de escala para o bloco intermediário.
  • downsample_padding (int, opcional, padrão 1) — O padding para a convolução de downsample.
  • downsample_type (str, opcional, padrão conv) — O tipo de downsample para as camadas de downsample. Escolha entre “conv” e “resnet”
  • upsample_type (str, opcional, padrão conv) — O tipo de upsample para as camadas de upsampling. Escolha entre “conv” e “resnet”
  • dropout (float, opcional, padrão 0.0) — A probabilidade de dropout a ser utilizada.
  • act_fn (str, opcional, padrão "silu") — A função de ativação a ser utilizada.
  • attention_head_dim (int, opcional, padrão 8) — A dimensão da cabeça de atenção.
  • norm_num_groups (int, opcional, padrão 32) — O número de grupos para normalização.
  • attn_norm_num_groups (int, opcional, padrão None) — Se definido como um inteiro, uma camada de group norm será criada na camada Attention do bloco intermediário com o número de grupos especificado. Se mantido como None, a camada de group norm será criada apenas se resnet_time_scale_shift estiver definido como default, e, se criada, terá norm_num_groups grupos.
  • norm_eps (float, opcional, padrão 1e-5) — O epsilon para normalização.
  • resnet_time_scale_shift (str, opcional, padrão "default") — Configuração de deslocamento de escala de tempo para blocos ResNet (veja ResnetBlock2D). Escolha entre default ou scale_shift.
  • class_embed_type (str, opcional, padrão None) — O tipo de incorporação de classe a ser usado, que é somado às incorporações de tempo. Escolha entre None, "timestep" ou "identity".
  • num_class_embeds (int, opcional, padrão None) — Dimensão de entrada da matriz de incorporação aprendível a ser projetada para time_embed_dim ao realizar condicionamento por classe com class_embed_type igual a None.

Um modelo UNet 2D que recebe uma amostra com ruído e um timestep e retorna uma saída com a forma da amostra.

Este modelo herda de ModelMixin. Consulte a documentação da superclasse para seus métodos genéricos implementados para todos os modelos (como download ou salvamento).

forward

< >

( sample: Tensortimestep: typing.Union[torch.Tensor, float, int]class_labels: typing.Optional[torch.Tensor] = Nonereturn_dict: bool = True ) → UNet2DOutput or tuple

Parâmetros

  • sample (torch.Tensor) — O tensor de entrada com ruído, com a seguinte forma (batch, channel, height, width).
  • timestep (torch.Tensor ou float ou int) — O número de timesteps para desruidar uma entrada.
  • class_labels (torch.Tensor, opcional, padrão None) — Rótulos de classe opcionais para condicionamento. Os embeddings deles serão somados aos embeddings de timestep.
  • return_dict (bool, opcional, padrão True) — Se deve ou não retornar um UNet2DOutput em vez de uma tupla simples.

Retorna

UNet2DOutput ou tuple

Se return_dict for True, um UNet2DOutput é retornado; caso contrário, uma tuple é retornada, em que o primeiro elemento é o tensor da amostra.

O método forward do UNet2DModel.

UNet2DOutput

classe diffusers.models.unets.unet_2d.UNet2DOutput

< >

( sample: Tensor )

Parâmetros

  • sample (torch.Tensor de forma (batch_size, num_channels, height, width)) — Os estados ocultos de saída da última camada do modelo.

A saída de UNet2DModel.

Update on GitHub

←UNet2DConditionModel UNet3DConditionModel→