stylegan2 pytorch

Simple stylegan2 para pytorch

Implementação simples de pytorch de stylegan2 com base em https://arxiv.org/abs/1912.04958 que pode ser completamente treinado a partir da linha de comando, não é necessária codificação.

Abaixo estão algumas flores que não existem.

Nem essas mãos

Nem essas cidades

Nem essas celebridades (treinadas por @yoniker)

Instalar

Você precisará de uma máquina com uma GPU e CUDA instalada. Então PIP instale o pacote como este

$ pip install stylegan2_pytorch

Se você estiver usando uma máquina Windows, os seguintes comandos funcionam.

$ conda install pytorch torchvision -c python
$ pip install stylegan2_pytorch

Usar

$ stylegan2_pytorch --data /path/to/images

É isso. As imagens de amostra serão salvas nos results/default e os modelos serão salvos periodicamente nos models/default .

Uso avançado

Você pode especificar o nome do seu projeto com

$ stylegan2_pytorch --data /path/to/images --name my-project-name

Você também pode especificar o local onde os resultados intermediários e os pontos de verificação do modelo devem ser armazenados com

$ stylegan2_pytorch --data /path/to/images --name my-project-name --results_dir /path/to/results/dir --models_dir /path/to/models/dir

Você pode aumentar a capacidade da rede (que padrão é 16 ) para melhorar os resultados da geração, ao custo de mais memória.

$ stylegan2_pytorch --data /path/to/images --network-capacity 256

Por padrão, se o treinamento for cortado, ele será retomado automaticamente do último arquivo apontado para verificação. Se você deseja reiniciar com novas configurações, basta adicionar uma new bandeira

$ stylegan2_pytorch --new --data /path/to/images --name my-project-name --image-size 512 --batch-size 1 --gradient-accumulate-every 16 --network-capacity 10

Depois de terminar o treinamento, você pode gerar imagens do seu último ponto de verificação como assim.

$ stylegan2_pytorch  --generate

Para gerar um vídeo de uma interpolação através de dois pontos aleatórios no espaço latente.

$ stylegan2_pytorch --generate-interpolation --interpolation-num-steps 100

Para salvar cada quadro individual da interpolação

$ stylegan2_pytorch --generate-interpolation --save-frames

Se um ponto de verificação anterior continha um gerador melhor (que geralmente acontece quando os geradores começam a se degradar no final do treinamento), você pode carregar de um ponto de verificação anterior com outra bandeira

$ stylegan2_pytorch --generate --load-from {checkpoint number}

Uma técnica usada no Stylegan e no Biggan está truncando os valores latentes, para que seus valores caam perto da média. Quanto pequeno o valor do truncamento, melhor as amostras aparecerão ao custo da variedade de amostras. Você pode controlar isso com o --trunc-psi , onde os valores normalmente caem entre 0.5 e 1 . Está definido em 0.75 como padrão

$ stylegan2_pytorch --generate --trunc-psi 0.5

Treinamento multi-GPU

Se você possui uma máquina com várias GPUs, o repositório oferece uma maneira de utilizar todos eles para treinamento. Com várias GPUs, cada lote será dividido uniformemente entre as GPUs disponíveis. Por exemplo, para 2 GPUs, com um tamanho de lotes de 32, cada GPU verá 16 amostras.

Você simplesmente precisa adicionar uma bandeira --multi-gpus , tudo é resolvido. Se você deseja restringir a GPUs específicas, pode usar a variável de ambiente CUDA_VISIBLE_DEVICES para controlar quais dispositivos podem ser usados. (Ex. CUDA_VISIBLE_DEVICES=0,2,3 Somente dispositivos 0, 2, 3 estão disponíveis)

$ stylegan2_pytorch --data ./data --multi-gpus --batch-size 32 --gradient-accumulate-every 1

Baixas quantidades de dados de treinamento

No passado, Gans precisava de muitos dados para aprender a gerar bem. O modelo FACES tirou 70 mil imagens de alta qualidade do Flickr, como exemplo.

No entanto, no mês de maio de 2020, pesquisadores de todo o mundo convergiram independentemente em uma técnica simples para reduzir esse número para 1-2k . Essa idéia simples era aumentar de maneira diferenciada todas as imagens, geradas ou reais, entrando no discriminador durante o treinamento.

Se alguém aumentar com uma probabilidade baixa o suficiente, os aumentos não 'vazarão' nas gerações.

Na configuração de dados baixos, você pode usar o recurso com um sinalizador simples.

 # find a suitable probability between 0. -> 0.7 at maximum
$ stylegan2_pytorch --data ./data --aug-prob 0.25

Por padrão, os aumentos utilizados são translation e cutout . Se você deseja adicionar color , pode fazê-lo com o argumento --aug-types .

 # make sure there are no spaces between items!
$ stylegan2_pytorch --data ./data --aug-prob 0.25 --aug-types [translation,cutout,color]

Você pode personalizá -lo para qualquer combinação dos três que desejar. O código de aumento diferenciável foi copiado e ligeiramente modificado daqui.

Quando faço para parar de treinar?

Durante o maior tempo possível, até o jogo adversário entre as duas redes neurais desmoronar (chamamos essa divergência). Por padrão, o número de etapas de treinamento é definido como 150000 para imagens de 128x128, mas você certamente deseja que esse número seja maior se o GaN não divergir até o final do treinamento ou se estiver treinando com uma resolução mais alta.

$ stylegan2_pytorch --data ./data --image-size 512 --num-train-steps 1000000

Atenção

Essa estrutura também permite que você adicione uma forma eficiente de auto-atimento às camadas designadas do discriminador (e a camada simétrica do gerador), o que melhorará bastante os resultados. Quanto mais atenção você puder pagar, melhor!

 # add self attention after the output of layer 1
$ stylegan2_pytorch - - data . / data - - attn - layers 1 

 # add self attention after the output of layers 1 and 2
# do not put a space after the comma in the list!
$ stylegan2_pytorch - - data . / data - - attn - layers [ 1 , 2 ]

Bônus

Treinamento em imagens transparentes

$ stylegan2_pytorch --data ./transparent/images/path --transparent

Considerações de memória

Quanto mais memória GPU você tiver, maior e melhor será a geração de imagens. A NVIDIA recomendou ter até 16 GB para o treinamento de imagens 1024x1024. Se você tiver menos do que isso, existem algumas configurações com as quais você pode brincar para que o modelo se encaixe.

$ stylegan2_pytorch --data /path/to/data 
    --batch-size 3 
    --gradient-accumulate-every 5 
    --network-capacity 16

1. Tamanho do lote-você pode diminuir o batch-size para 1, mas deve aumentar o gradient-accumulate-every de correspondência para que o mini-lote que a rede vê não é muito pequeno. Isso pode ser confuso para um leigo, então vou pensar em como automatizaria a escolha do gradient-accumulate-every os quais estão no futuro.

2. Capacidade da rede - você pode diminuir a capacidade da rede neural para diminuir os requisitos de memória. Esteja ciente de que isso demonstrou degradar o desempenho da geração.

Se nada disso funcionar, você poderá se contentar com o GaN 'leve', o que permitirá que você compense a qualidade do treino com maiores resoluções em quantidade razoável de tempo.

Implantação na AWS

Abaixo estão algumas etapas que podem ser úteis para a implantação usando os serviços da Web da Amazon. Para usar isso, você precisará provisionar uma instância EC2 apoiada por GPU. Um tipo de instância apropriado seria de uma série P2 ou P3. Eu (iBoates) tentei um p2.xlarge (a opção mais barata) e era bastante lenta, mais lenta do que usar o Google Colab. Tipos de instância mais poderosos podem ser melhores, mas são mais caros. Você pode ler mais sobre eles aqui.

Etapas de configuração

1. Arquive seus dados de treinamento e envie -os para um balde S3
2. Provisionar sua instância do EC2 (usei um Ubuntu ami)
3. Faça login na sua instância EC2 via SSH
4. Instale o cliente da AWS CLI e configure -o:

sudo snap install aws-cli --classic
aws configure

Você precisará inserir suas chaves de acesso da AWS, que você pode recuperar do console de gerenciamento sob o console de gerenciamento da AWS> perfil> minhas credenciais de segurança> chaves de acesso

Em seguida, execute esses comandos, ou talvez os coloque em um script de shell e execute isso:

mkdir data
curl -O https://bootstrap.pypa.io/get-pip.py
sudo apt-get install python3-distutils
python3 get-pip.py
pip3 install stylegan2_pytorch
export PATH= $PATH :/home/ubuntu/.local/bin
aws s3 sync s3:// < Your bucket name > ~ /data
cd data
tar -xf ../train.tar.gz

Agora você deve ser capaz de treinar simplesmente chamando stylegan2_pytorch [args] .

Notas:

• Se você tiver muitos dados de treinamento, pode ser necessário provisionar o armazenamento extra de bloco via EBS.
• Além disso, pode ser necessário espalhar seus dados por vários arquivos.
• Você deve executar isso na janela de uma screen para que ele não termine quando você sair da sessão SSH.

Pesquisar

Pontuações FID

Graças ao getSeclético, agora você pode calcular a pontuação do FID periodicamente! Novamente, fez super simples com um argumento extra, como mostrado abaixo.

Em primeiro lugar, instale o pacote pytorch_fid

$ pip install pytorch-fid

Seguido pela

$ stylegan2_pytorch --data ./data --calculate-fid-every 5000

Os resultados do FID serão registrados para ./results/{name}/fid_scores.txt name }/fid_scores.txt

Codificação

Se você deseja amostrar imagens programaticamente, pode fazê -lo com a seguinte classe ModelLoader .

 import torch
from torchvision . utils import save_image
from stylegan2_pytorch import ModelLoader

loader = ModelLoader (
    base_dir = '/path/to/directory' ,   # path to where you invoked the command line tool
    name = 'default'                   # the project name, defaults to 'default'
)

noise   = torch . randn ( 1 , 512 ). cuda () # noise
styles  = loader . noise_to_styles ( noise , trunc_psi = 0.7 )  # pass through mapping network
images  = loader . styles_to_images ( styles ) # call the generator on intermediate style vectors

save_image ( images , './sample.jpg' ) # save your images, or do whatever you desire

Log para rastreador de experimento

Para registrar as perdas a um rastreador de experimento de código aberto (AIM), basta passar por um sinalizador extra assim.

$ stylegan2_pytorch --data ./data --log

Em seguida, você precisa ter certeza de que está instalado no Docker. Seguindo as instruções no AIM, você executa o seguinte no seu terminal.

$ aim up

Em seguida, abra seu navegador para o endereço e você deve ver

Experimental

Treinamento principal para gerador

Um novo artigo produziu evidências de que, simplesmente zero nas contribuições de gradiente de amostras que são consideradas falsas pelo discriminador, o gerador aprende significativamente melhor, alcançando um novo estado da arte.

$ stylegan2_pytorch - - data . / data - - top - k - training

Gamma é um cronograma de decaimento que diminui lentamente o Topk do tamanho completo do lote para a fração alvo de 50% (também hiperparâmetro modificável).

$ stylegan2_pytorch - - data . / data - - top - k - training - - generate - top - k - frac 0.5 - - generate - top - k - gamma 0.99

Quantização de recursos

Um artigo recente relatou melhores resultados se as representações intermediárias do discriminador forem quantizadas vetoriais. Embora eu não tenha notado nenhuma alteração dramática, decidi adicionar isso como um recurso, para que outras mentes por aí possam investigar. Para usar, você deve especificar qual camada (s) você gostaria de vetorizar quantizar. O tamanho do dicionário padrão é 256 e também é ajustável.

 # feature quantize layers 1 and 2, with a dictionary size of 512 each
# do not put a space after the comma in the list!
$ stylegan2_pytorch - - data . / data - - fq - layers [ 1 , 2 ] - - fq - dict - size 512

Regularização de perdas contrastivas

Eu tentei aprendizado contrastante sobre o discriminador (em passo com o treinamento habitual de GaN) e possivelmente observei melhoria a estabilidade e a qualidade dos resultados finais. Você pode ativar esse recurso experimental com um sinalizador simples, como mostrado abaixo.

$ stylegan2_pytorch - - data . / data - - cl - reg

Perda de discriminador relativista

Isso foi proposto no papel gan relativista para estabilizar o treinamento. Eu tive resultados mistos, mas incluirei o recurso para quem deseja experimentar.

$ stylegan2_pytorch - - data . / data - - rel - disc - loss

Bloco 4x4 não constante

Por padrão, a arquitetura Stylegan estiliza um bloco 4x4 constante, à medida que é progressivamente amostrado. Esse é um recurso experimental que o torna para que o bloco 4x4 seja aprendido com o vetor de estilo w em vez disso.

$ stylegan2_pytorch - - data . / data - - no - const

Perda contrastiva dupla

Um artigo recente propôs que uma nova perda contrastiva entre os logits reais e falsos possa melhorar a qualidade em relação a outros tipos de perdas. (O padrão neste repositório é a perda de dobradiça e o artigo mostra uma ligeira melhoria)

$ stylegan2_pytorch - - data . / data - - dual - contrast - loss 

Alternativas

Stylegan2 + discriminador unet

Eu obtive resultados muito bons com um discriminador não -ET, mas a mudança arquitetonicamente foi grande demais para se encaixar como uma opção neste repositório. Se você está buscando a perfeição, fique à vontade para experimentá -lo.

Se você gostaria que eu desse o tratamento real a outra arquitetura GaN (Biggan), fique à vontade para chegar ao meu e -mail. Fico feliz em ouvir seu arremesso.

Apreciação

Obrigado a Matthew Mann por seu inspirador porto simples para o tensorflow 2.0

Referências

 @article { Karras2019stylegan2 ,
    title   = { Analyzing and Improving the Image Quality of {StyleGAN} } ,
    author  = { Tero Karras and Samuli Laine and Miika Aittala and Janne Hellsten and Jaakko Lehtinen and Timo Aila } ,
    journal = { CoRR } ,
    volume  = { abs/1912.04958 } ,
    year    = { 2019 } ,
}

 @misc { zhao2020feature ,
    title   = { Feature Quantization Improves GAN Training } ,
    author  = { Yang Zhao and Chunyuan Li and Ping Yu and Jianfeng Gao and Changyou Chen } ,
    year    = { 2020 }
}

 @misc { chen2020simple ,
    title   = { A Simple Framework for Contrastive Learning of Visual Representations } ,
    author  = { Ting Chen and Simon Kornblith and Mohammad Norouzi and Geoffrey Hinton } ,
    year    = { 2020 }
}

 @article {,
    title     = { Oxford 102 Flowers } ,
    author    = { Nilsback, M-E. and Zisserman, A., 2008 } ,
    abstract  = { A 102 category dataset consisting of 102 flower categories, commonly occuring in the United Kingdom. Each class consists of 40 to 258 images. The images have large scale, pose and light variations. }
}

 @article { afifi201911k ,
    title   = { 11K Hands: gender recognition and biometric identification using a large dataset of hand images } ,
    author  = { Afifi, Mahmoud } ,
    journal = { Multimedia Tools and Applications }
}

 @misc { zhang2018selfattention ,
    title   = { Self-Attention Generative Adversarial Networks } ,
    author  = { Han Zhang and Ian Goodfellow and Dimitris Metaxas and Augustus Odena } ,
    year    = { 2018 } ,
    eprint  = { 1805.08318 } ,
    archivePrefix = { arXiv }
}

 @article { shen2019efficient ,
    author    = { Zhuoran Shen and
               Mingyuan Zhang and
               Haiyu Zhao and
               Shuai Yi and
               Hongsheng Li } ,
    title     = { Efficient Attention: Attention with Linear Complexities } ,
    journal   = { CoRR } ,  
    year      = { 2018 } ,
    url       = { http://arxiv.org/abs/1812.01243 } ,
}

 @article { zhao2020diffaugment ,
    title   = { Differentiable Augmentation for Data-Efficient GAN Training } ,
    author  = { Zhao, Shengyu and Liu, Zhijian and Lin, Ji and Zhu, Jun-Yan and Han, Song } ,
    journal = { arXiv preprint arXiv:2006.10738 } ,
    year    = { 2020 }
}

 @misc { zhao2020image ,
    title  = { Image Augmentations for GAN Training } ,
    author = { Zhengli Zhao and Zizhao Zhang and Ting Chen and Sameer Singh and Han Zhang } ,
    year   = { 2020 } ,
    eprint = { 2006.02595 } ,
    archivePrefix = { arXiv }
}

 @misc { karras2020training ,
    title   = { Training Generative Adversarial Networks with Limited Data } ,
    author  = { Tero Karras and Miika Aittala and Janne Hellsten and Samuli Laine and Jaakko Lehtinen and Timo Aila } ,
    year    = { 2020 } ,
    eprint  = { 2006.06676 } ,
    archivePrefix = { arXiv } ,
    primaryClass = { cs.CV }
}

 @misc { jolicoeurmartineau2018relativistic ,
    title   = { The relativistic discriminator: a key element missing from standard GAN } ,
    author  = { Alexia Jolicoeur-Martineau } ,
    year    = { 2018 } ,
    eprint  = { 1807.00734 } ,
    archivePrefix = { arXiv } ,
    primaryClass = { cs.LG }
}

 @misc { sinha2020topk ,
    title   = { Top-k Training of GANs: Improving GAN Performance by Throwing Away Bad Samples } ,
    author  = { Samarth Sinha and Zhengli Zhao and Anirudh Goyal and Colin Raffel and Augustus Odena } ,
    year    = { 2020 } ,
    eprint  = { 2002.06224 } ,
    archivePrefix = { arXiv } ,
    primaryClass = { stat.ML }
}

 @misc { yu2021dual ,
    title   = { Dual Contrastive Loss and Attention for GANs } ,
    author  = { Ning Yu and Guilin Liu and Aysegul Dundar and Andrew Tao and Bryan Catanzaro and Larry Davis and Mario Fritz } ,
    year    = { 2021 } ,
    eprint  = { 2103.16748 } ,
    archivePrefix = { arXiv } ,
    primaryClass = { cs.CV }
}