VQ VAE on MNIST

AutoEncoder variacional-intermitido por vetor (VQ-VAE)

O repositório consiste em um VQ-VAE implementado em Pytorch e treinado no conjunto de dados MNIST.

VQ-VAE: Visão geral

O VQ-VAE segue o mesmo conceito básico que por trás dos codificadores automáticos variacionais (VAE). O VQ-VAE usa incorporações latentes discretas para codificadores automáticos variacionais , ou seja, cada dimensão do z (vetor latente) é um número inteiro discreto, em vez da distribuição normal contínua geralmente usada durante a codificação das entradas.

Vaes consistem em 3 partes:

1. Uma rede de codificadores que parametra o q posterior (z | x) sobre os latentes
2. Uma distribuição anterior p (z)
3. Um decodificador com distribuição p (x | z) sobre dados de entrada

Bem, você pode perguntar sobre as diferenças que o VQ-VAES traz para a mesa. Vamos listá -los:

1. Os codificadores modelam uma distribuição categórica, amostragem da qual você obtém valores integrais
2. Esses valores integrais são usados ​​para indexar um dicionário de incorporação
3. Os valores indexados são então passados ​​para o decodificador

Por que apresentar as diferenças?

Muitos objetos importantes do mundo real são discretos. Por exemplo, em imagens, podemos ter categorias como "gato", "carro" etc. e pode não fazer sentido interpolar entre essas categorias. Representações discretas também são mais fáceis de modelar.

Arquitetura

onde:

• n : tamanho do lote
• h : Altura da imagem
• w : largura da imagem
• c : Número de canais na imagem de entrada
• d : Número de canais no estado oculto

Trabalhando

Aqui está uma breve visão geral do funcionamento de uma rede VQ-VAE:

1. O VQ-VAE consiste em um codificador, uma incorporação (ou um livro de códigos) e um decodificador.
2. Quando uma imagem é passada como entrada, ela é convertida em vetores latentes usando a rede de codificadores .

3. O espaço de incorporação consiste em muitos vetores latentes, que são comparados ao da entrada.
4. As distâncias são calculadas e o vetor latente mais semelhante (menor de distância) (no espaço de incorporação) ao vetor latente da entrada é selecionado.
5. O selecionado é alimentado na rede de decodificadores que reconstrói a imagem .

Camada de quantização de vetores

O funcionamento da camada VQ pode ser explicado em seis etapas, conforme numerado na figura:

1. Remodapa: Todas
2. Cálculo de distâncias: para cada um dos vetores n h w, calculamos a distância de cada um dos k vetores do dicionário de incorporação para obter uma matriz de forma (n h w, k)
3. Argmin: Para cada um dos vetores n h w, encontramos o índice dos vetores mais próximos dos k do dicionário
4. ÍNDICE DO DICIONÁRIO: INDEX O vetor mais próximo do dicionário para cada um dos vetores n h w
5. Remodapa: converta de volta à forma (n, h, w, d)
6. Copiando gradientes: não é possível treinar essa arquitetura através da propagação, pois o gradiente não flui através de Argmin. Portanto, tentamos aproximar copiando os gradientes de z_q de volta para z_e. Dessa forma, não estamos na verdade minimizando a função de perda, mas ainda é capaz de transmitir algumas informações para treinamento.

Funções de perda

O VQ-VAE usa 3 perdas para calcular a perda total durante o treinamento:

1. Perda de reconstrução: otimiza o decodificador e o codificador como VAE, ou seja, a diferença entre a imagem de entrada e a reconstrução:

   reconstruction_loss = -log( p(x|z_q) )

2. Perda do livro de código: devido ao fato de os gradientes ignorarem a incorporação, um algoritmo de aprendizado de dicionário que usa um erro L2 para mover os vetores de incorporação E_I para a saída do codificador é usada.

   codebook_loss = ‖ sg[z_e(x)]− e ‖^2

   (SG representa o operador de gradiente de parada significa que nenhum gradiente flui através do que quer que seja aplicado)

3. Perda de compromisso: Como o volume do espaço de incorporação é adimensional, ele pode crescer arbitrariamente se as incorporações E_I não treinarem tão rápido quanto os parâmetros do codificador e, portanto, uma perda de compromisso é adicionada para garantir que o codificador se comprometa a uma incorporação.

   commitment_loss = β‖ z_e(x)− sg[e] ‖^2

   (β é um hiperparâmetro que controla o quanto queremos pesar a perda de compromisso em comparação com outros componentes)

Conteúdo

1. Instruções de configuração
2. Treinando seu modelo do zero
3. Gerando imagens do modelo
4. Visão geral do repositório
5. Resultados
   1. Treinando imagens
   2. Gráficos de treinamento
   3. Testando gráficos
   4. Imagens geradas
6. Observações
7. Créditos

1. Instruções de configuração

Você pode baixar o repo ou cloná -lo executando o seguinte no prompt de cmd

 https://github.com/praeclarumjj3/VQ-VAE-on-MNIST.git

2. Treinando seu modelo do zero

Você pode treinar o modelo do zero pelo seguinte comando (no Google Colab)

 ! python3 VQ-VAE.py --output-folder [NAME_OF_OUTPUT_FOLDER] --data-folder [PATH_TO_MNIST_dataset] --device ['cpu' or 'cuda' ] --hidden-size [SIZE] --k [NUMBER] --batch-size [BATCH_SIZE] --num_epoch [NUMBER_OF_EPOCHS] --lr [LEARNING_RATE] --beta [VALUE] --num-workers [NUMBER_OF_WORKERS]

• output-folder - Nome da pasta de dados
• data-folder - Nome da pasta de dados
• device - Defina o dispositivo (CPU ou CUDA, padrão: CPU)
• hidden-size - tamanho dos vetores latentes (padrão: 40)
• k - Número de vetores latentes (Padrão: 512)
• batch-size - tamanho do lote (padrão: 128)
• num-epochs - Número de épocas (Padrão: 10)
• lr - Taxa de aprendizado para Adam Optimizer (Padrão: 2E -4)
• beta - Contribuição da perda de compromisso, entre 0,1 e 2,0 (padrão: 1,0)
• num-workers - Número de trabalhadores para trajetórias amostragem (Padrão: CPU_COUNT () - 1)

O programa baixa automaticamente o conjunto de dados MNIST e o salva na pasta PATH_TO_MNIST_dataset (você precisa criar esta pasta). Isso só acontece uma vez.

Ele também cria uma pasta logs e pasta models e dentro deles cria uma pasta com o nome passado por você para salvar logs e modelar pontos de verificação dentro dele, respectivamente.

3. Gerando imagens do modelo

Para gerar novas imagens de Z amostradas aleatoriamente a partir de uma unidade Gaussian Run o seguinte comando (no Google Colab):

 ! python3 generate.py  --model [SAVED_MODEL_FILENAME] --input [MNIST_or_random] --device ['cpu' or 'cuda' ] --hidden-size [SIZE] --k [NUMBER] --filename [SAVING_NAME] 

• model - nome do arquivo contendo o modelo
• input - mnist ou aleatório
• device - Defina o dispositivo (CPU ou CUDA, padrão: CPU)
• hidden-size - tamanho dos vetores latentes (padrão: 40)
• k - Número de vetores latentes (Padrão: 512)
• filename - nome com qual arquivo deve ser salvo

Ele gera uma grade de 10*10 de imagens que são salvas em uma pasta chamada generatedImages .

Você pode usar um modelo pré-treinado baixando-o no link no model.txt .

4. Visão geral do repositório

O repositório contém os seguintes arquivos

• modules.py - contém os diferentes módulos usados ​​para fazer nosso modelo
• VQ-VAE.py -contém as funções e o código para treinar nosso modelo VQ-VAE
• vector_quantizer.py - As classes de quantização do vetor são definidas neste arquivo
• generate-py -gera novas imagens de um modelo pré-treinado
• model.txt - contém um link para um modelo pré -treinado
• README.md - Readme dando uma visão geral do repositório
• references.txt - Referências usadas ao criar este repositório
• readme_images - tem várias imagens para o ReadMe
• MNIST - contém o conjunto de dados MNIST com zíper (embora ele seja baixado automaticamente, se necessário)
• Training track for VQ-VAE.txt -Contém os valores de perda durante o treinamento do nosso modelo VQ-VAE
• logs_VQ-VAE -contém os logs de tensorboard zipped para o nosso modelo VQ-VAE (criado automaticamente pelo programa)
• testers.py - contém algumas funções para testar nossos módulos definidos

Comando para executar o Tensorboard (no Google Colab):

 %load_ext tensorboard
%tensordboard --logdir [path_to_folder_with_logs]

5. Resultados

1. Treinando imagens

Imagem de treinamento

Imagem de 0th época

Imagem da 2ª época

Imagem da 4ª época

Imagem da 6ª época

Imagem da 8ª época

Imagem da 10ª época

As reconstruções continuam melhorando e, no final, quase se assemelham às imagens do treinamento_set que se reflete nos valores de perda (verifique a Training track for VQ-VAE.txt ).

2. Gráficos de treinamento

Perda de reconstrução

Perda de quantização

Total_loss

A perda total, a perda de reconstrução e a perda de quantização diminuem uniformemente conforme o esperado.

3. Testando gráficos

Testing_loss

A perda de teste diminui uniformemente conforme o esperado.

4. Imagens geradas

A grade de imagem a seguir foi gerada após a passagem de imagens MNIST como entradas:

A geração é muito boa.

As seguintes grades de imagem foram geradas após a passagem de AZ amostrada aleatoriamente de uma unidade gaussiana como entrada para o modelo e depois passou pelo decodificador

As imagens não parecem perfeitas. Ajustando as dimensões do espaço latente, o número de vetores de incorporação etc. pode ajudar a gerar melhores imagens aleatórias.

6. Observações

O modelo foi treinado no Google Colab por 10 épocas, com tamanho de lote 128.

Depois de treinar, o modelo conseguiu reconstruir muito bem as imagens de entrada e também conseguiu gerar novas imagens, embora as imagens geradas não sejam tão boas.
O treinamento e a perda de testes também continuaram diminuindo quase monotonicamente.

Observei que o treinamento do modelo para mais de 10 a 20 épocas produziu resultados que sugeriam um provável sinal de excesso de ajuste no modelo. Além disso, experimentei diferentes dimensões do espaço lado e, na dimension = 40 produziram os melhores resultados. O melhor alcance para a dimensão chegou entre 16-42.

7. Créditos

As fontes a seguir ajudaram muito a tornar este repositório

• Learning de representação discreta neural - Aaron van den Oord, Oriol Vinyals, Koray Kavukcuoglu
• Gerando diversas imagens de alta fidelidade com VQ-VAE-2-Ali Razavi, Aaron Van Den Oord, Oriol Vinyals
• https://nbviewer.jupyter.org/github/zalandoresearch/pytorch-vq-vae/blob/master/vq-vae.ipynb
• https://www.kaggle.com/ameroyer/keras-vq-vae-for-image-generation
• https://blog.usejournal.com/understanding-vector-quantized-variational-autoencoders-vq-vae-323d710a888a
• https://christineai.blog/pixelcnn-and-pixelrnn/
• https://github.com/ritheshkumar95/pytorch-vqvae
• https://github.com/ayushtues/genzoo