pytorch pretrained BigGAN

Pytorch pré -criou o biggan

Uma reimplementação de Pytorch Op-for-Op do modelo Biggan de DeepMind com os pesos pré-treinados de DeepMind.

Introdução

Este repositório contém uma reimplementação de Pytorch Op-for-Op do Biggan de Deepmind, que foi lançado com o treinamento GaN de grande escala para a síntese de imagem natural de alta fidelidade por Andrew Brock, Jeff Donahue e Karen Simonyan.

Essa implementação de Pytorch do Biggan é fornecida com os modelos de 128x128, 256x256 e 512x512 da DeepMind. Também fornecemos os scripts usados para baixar e converter esses modelos nos modelos TensorFlow Hub.

Essa reimplementação foi feita a partir do gráfico de computação bruta da versão Tensorflow e se comporta de maneira semelhante à versão do tensorflow (variação da diferença de saída da ordem de 1E-5).

Atualmente, essa implementação contém apenas o gerador, pois os pesos do discriminador não foram liberados (embora a estrutura do discriminador seja muito semelhante ao gerador para que possa ser adicionado com bastante facilidade. Diga -me se você deseja fazer um PR nisso, ficaria feliz em ajudar.)

Instalação

Este repo foi testado no Python 3.6 e Pytorch 1.0.1

Pytorch Biggan pré -traido pode ser instalado a partir da PIP da seguinte forma:

pip install pytorch-pretrained-biggan

Se você simplesmente quer brincar com o GAN, isso deve ser suficiente.

Se você deseja usar os scripts de conversão e os utilitários do ImageNet, são necessários requisitos adicionais, em particular o tensorflow e o NLTK. Para instalar todos os requisitos, use o arquivo full_requirements.txt :

git clone https://github.com/huggingface/pytorch-pretrained-BigGAN.git
cd pytorch-pretrained-BigGAN
pip install -r full_requirements.txt

Modelos

Este repositório fornece acesso direto e simples às versões "profundas" pré -treinadas do Biggan para resoluções de 128, 256 e 512 pixels, conforme descrito na publicação associada. Aqui estão alguns detalhes sobre os modelos:

• BigGAN-deep-128 : um modelo de parâmetros de 50,4m gerando imagens de 128x128 pixels, o modelo de despejo de 201 MB,
• BigGAN-deep-256 : um modelo de 55,9m de parâmetros gerando imagens de 256x256 pixels, o modelo de despejo de 224 MB,
• BigGAN-deep-512 : Um modelo de 56,2m de parâmetros que gera imagens de 512x512 pixels, o modelo de despejo de 225 Mb.

Consulte o Apêndice B do artigo para obter detalhes sobre as arquiteturas.

Todos os modelos compreendem estatísticas de norma em lote pré-computadas para 51 valores de truncamento entre 0 e 1 (consulte o Apêndice C.1 no artigo para obter detalhes).

Uso

Aqui está um exemplo de partida rápida usando BigGAN com um modelo pré-treinado.

Consulte a seção Doc abaixo para obter detalhes sobre essas classes e métodos.

 import torch
from pytorch_pretrained_biggan import ( BigGAN , one_hot_from_names , truncated_noise_sample ,
                                       save_as_images , display_in_terminal )

# OPTIONAL: if you want to have more information on what's happening, activate the logger as follows
import logging
logging . basicConfig ( level = logging . INFO )

# Load pre-trained model tokenizer (vocabulary)
model = BigGAN . from_pretrained ( 'biggan-deep-256' )

# Prepare a input
truncation = 0.4
class_vector = one_hot_from_names ([ 'soap bubble' , 'coffee' , 'mushroom' ], batch_size = 3 )
noise_vector = truncated_noise_sample ( truncation = truncation , batch_size = 3 )

# All in tensors
noise_vector = torch . from_numpy ( noise_vector )
class_vector = torch . from_numpy ( class_vector )

# If you have a GPU, put everything on cuda
noise_vector = noise_vector . to ( 'cuda' )
class_vector = class_vector . to ( 'cuda' )
model . to ( 'cuda' )

# Generate an image
with torch . no_grad ():
    output = model ( noise_vector , class_vector , truncation )

# If you have a GPU put back on CPU
output = output . to ( 'cpu' )

# If you have a sixtel compatible terminal you can display the images in the terminal
# (see https://github.com/saitoha/libsixel for details)
display_in_terminal ( output )

# Save results as png images
save_as_images ( output )

Doc

Carregando pesos pré-treinados do DeepMind

Para carregar um dos modelos pré-treinados do DeepMind, instancie um modelo BigGAN com from_pretrained() como:

 model = BigGAN . from_pretrained ( PRE_TRAINED_MODEL_NAME_OR_PATH , cache_dir = None )

onde

• PRE_TRAINED_MODEL_NAME_OR_PATH é:

  • O nome de atalho de um modelo pré-treinado do Google AI ou OpenAI selecionado na lista:

    • biggan-deep-128 : 12 camadas, 768 parâmetros de 12 cabeças, 110m
    • biggan-deep-256 : parâmetros de 24 camadas, 1024 ocultos, 16 cabeças, 340m
    • biggan-deep-512 : parâmetros de 12 camadas, 768 ocultos, 12 cabeças e 110m

  • Um caminho ou URL para um arquivo de modelo pré -terenciado contendo:

    • config.json : um arquivo de configuração para o modelo e
    • pytorch_model.bin Um dump pytorch de uma instância pré-treinada de BigGAN (salva com a torch.save() ).

  Se PRE_TRAINED_MODEL_NAME_OR_PATH for um nome de atalho, os pesos pré-treinados serão baixados do AWS S3 (consulte os links aqui) e armazenados em uma pasta de cache para evitar downloads futuros (a pasta de cache poderá ser encontrada em ~/.pytorch_pretrained_biggan/ ).

• cache_dir pode ser um caminho opcional para um diretório específico para baixar e cache os pesos do modelo pré-treinado.

Configuração

BigGANConfig é uma classe para armazenar e carregar configurações do Biggan. Está definido em config.py .

Aqui estão alguns detalhes sobre os atributos:

• output_dim : resolução de saída do GaN (128, 256 ou 512) para os modelos pré-treinados,
• z_dim : Tamanho do vetor de ruído (128 para os modelos pré-treinados).
• class_embed_dim : tamanho dos vetores de incorporação de classe (128 para os modelos pré-treinados).
• channel_width : tamanho de cada canal (128 para os modelos pré-treinados).
• num_classes : número de classes no conjunto de dados de treinamento, como o ImageNet (1000 para os modelos pré-treinados).
• layers : uma lista de definição de camadas. Cada definição para uma camada é um triplo de [amostra superior na camada? (bool), número de canais de entrada (int), número de canais de saída (int)]
• attention_layer_position : Posição da camada de auto-ataque na hierarquia da camada (8 para os modelos pré-treinados).
• eps : Valor Epsilon a ser usado para camadas de normalização espectral e em lote (1E-4 para os modelos pré-treinados).
• n_stats : Número de estatísticas pré-computadas para as camadas de normalização do lote associadas a vários valores de truncamento entre 0 e 1 (51 para os modelos pré-treinados).

Modelo

BigGAN é um modelo Pytorch ( torch.nn.Module ) do biggan definido em model.py . Este modelo compreende as incorporações de classe (uma camada linear) e o gerador com uma série de convoluções e normas de lote condicional. Atualmente, o discriminador não é implementado, pois os pesos pré-treinados não foram liberados para ele.

As entradas e a saída são idênticas às entradas e saídas do modelo do tensorflow .

Nós os detalhamos aqui.

BigGAN toma como entrada :

• z : uma tocha.floattensor de forma [batch_size, config.z_dim] com ruído amostrado de uma distribuição normal truncada e
• class_label : uma tocha opcional.longtensor de forma [batch_size, sequence_length] com os índices dos tipos de token selecionados em [0, 1]. O tipo 0 corresponde a uma sentence A e tipo 1 corresponde a um token sentence B (consulte o papel Bert para obter mais detalhes).
• truncation : um flutuador entre 0 (não compreendido) e 1. A truncamento do normal truncado usado para criar o vetor de ruído. Esse valor de truncamento é usado para selecionar entre um conjunto de estatísticas pré-computadas (médias e variações) para as camadas de norma em lote.

BigGAN produz uma matriz de forma [Batch_size, 3, Resolução, Resolução] onde a resolução é 128, 256 ou 512, dependendo do modelo:

Utilitários: imagens, ruído, classes ImageNet

Fornecemos alguns métodos de utilidade para usar o modelo. Eles são definidos em utils.py .

Aqui estão alguns detalhes sobre esses métodos:

• truncated_noise_sample(batch_size=1, dim_z=128, truncation=1., seed=None) :

  Crie um vetor de ruído truncado.

  • Params:
    • Batch_size: tamanho do lote.
    • dim_z: dimensão de z
    • truncamento: valor de truncamento para usar
    • Semente: semente para o gerador aleatório
  • Saída: Matriz de forma (batch_size, dim_z)

• convert_to_images(obj) :

  Converta um tensor de saída do Biggan em uma lista de imagens.

  • Params:
    • OBJ: Tensor ou Numpy Grey of Shape (Batch_size, canais, altura, largura)
  • Saída:
    • Lista de imagens de travesseiro de tamanho (altura, largura)

• save_as_images(obj, file_name='output') :

  Converta e salve um tensor de saída do Biggan em uma lista de imagens salvas.

  • Params:
    • OBJ: Tensor ou Numpy Grey of Shape (Batch_size, canais, altura, largura)
    • File_name: Path e implorando o nome do arquivo para salvar. As imagens serão salvas como file_name_{image_number}.png

• display_in_terminal(obj) :

  Converta e exiba um tensor de saída do biggan no terminal. Esta função usa libsixel e só funcionará em um terminal compatível com LBsixel. Consulte https://github.com/saitoha/libsixel para obter mais detalhes.

  • Params:
    • OBJ: Tensor ou Numpy Grey of Shape (Batch_size, canais, altura, largura)
    • File_name: Path e implorando o nome do arquivo para salvar. As imagens serão salvas como file_name_{image_number}.png

• one_hot_from_int(int_or_list, batch_size=1) :

  Crie um vetor de um hot a partir de um índice de classe ou uma lista de índices de classe.

  • Params:
    • int_or_list: int, ou lista de int, das classes imagenet (entre 0 e 999)
    • Batch_size: tamanho do lote.
      • Se int_or_list for um int, crie um lote de classes idênticas.
      • Se int_or_list for uma lista, devemos ter len(int_or_list) == batch_size
  • Saída:
    • Matriz de forma (Batch_size, 1000)

• one_hot_from_names(class_name, batch_size=1) :

  Crie um vetor de um hot a partir do nome de uma classe ImageNet ('Tennis Ball', 'Daisy', ...). Usamos a pesquisa do WordNet do NLTK para tentar encontrar o Synset relevante do ImageNet e pegar o primeiro. Se não conseguirmos encontrá -lo de maneira diretamente, olhamos para os hipônimos e hipernyms do nome da classe.

  • Params:
    • Class_name: String contendo o nome de um objeto ImageNet.
  • Saída:
    • Matriz de forma (Batch_size, 1000)

Download e scripts de conversão

Os scripts para baixar e converter os modelos TensorFlow no Tensorflow Hub são fornecidos em ./scripts.

Os scripts podem ser usados diretamente como:

./scripts/download_tf_hub_models.sh
./scripts/convert_tf_hub_models.sh