Hugging Face's logo

Diffusers documentation

UNet2DModel

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→