VQ VAE on MNIST

Autoencoder variasional vektor (VQ-VAE)

Repositori terdiri dari VQ-VAE yang diimplementasikan di Pytorch dan dilatih pada dataset MNIST.

VQ-VAE: Ikhtisar

VQ-VAE mengikuti konsep dasar yang sama seperti di balik variasional auto-encoder (VAE). VQ-VAE menggunakan embeddings laten diskrit untuk encoder otomatis variasional , yaitu setiap dimensi z (vektor laten) adalah bilangan bulat diskrit, alih-alih distribusi normal kontinu yang umumnya digunakan saat mengkode input.

VAE terdiri dari 3 bagian:

1. Jaringan enkoder yang parametrisasi q posterior (z | x) di atas laten
2. Distribusi sebelumnya p (z)
3. Decoder dengan distribusi p (x | z) di atas data input

Nah, Anda mungkin bertanya tentang perbedaan yang dibawa VQ-VAE ke meja. Mari kita daftarkan mereka:

1. Encoders memodelkan distribusi kategorikal, pengambilan sampel dari mana Anda mendapatkan nilai integral
2. Nilai -nilai integral ini digunakan untuk mengindeks kamus embeddings
3. Nilai yang diindeks kemudian diteruskan ke decoder

Mengapa memperkenalkan perbedaannya?

Banyak objek dunia nyata yang penting diskrit. Misalnya dalam gambar kita mungkin memiliki kategori seperti "kucing", "mobil", dll. Dan mungkin tidak masuk akal untuk menginterpolasi antara kategori -kategori ini. Representasi diskrit juga lebih mudah dimodelkan.

Arsitektur

Di mana:

• n : Ukuran Batch
• h : Tinggi gambar
• w : Lebar gambar
• c : Jumlah saluran dalam gambar input
• d : Jumlah saluran dalam keadaan tersembunyi

Bekerja

Berikut gambaran singkat tentang kerja jaringan VQ-VAE:

1. VQ-VAE terdiri dari encoder, embedding (atau codebook) dan decoder.
2. Ketika gambar dilewatkan sebagai input, itu dikonversi menjadi vektor laten menggunakan jaringan encoder .

3. Ruang embedding terdiri dari banyak vektor laten, yang dibandingkan dengan yang dari input.
4. Jarak dihitung dan vektor laten yang paling mirip (paling tidak jarak) (dalam ruang embedding) ke vektor laten input dipilih.
5. Yang dipilih dimasukkan ke dalam jaringan dekoder yang merekonstruksi gambar .

Lapisan kuantisasi vektor

Kerja lapisan VQ dapat dijelaskan dalam enam langkah yang diberi nomor dalam gambar:

1. Reshape: Semua dimensi kecuali yang terakhir digabungkan menjadi satu sehingga kami memiliki vektor masing -masing dari dimensi d
2. Menghitung jarak: Untuk masing -masing vektor N H W kami menghitung jarak dari masing -masing vektor K dari kamus embedding untuk mendapatkan matriks bentuk (N H W, K)
3. Argmin: Untuk masing -masing vektor N H W kami menemukan indeks vektor terdekat dari kamus
4. Indeks dari kamus: indeks vektor terdekat dari kamus untuk masing -masing vektor n h w
5. Reshape: Konversi kembali ke bentuk (n, h, w, d)
6. Menyalin gradien: Tidak mungkin untuk melatih arsitektur ini melalui backpropagation karena gradien tidak akan mengalir melalui argmin. Oleh karena itu kami mencoba mendekati dengan menyalin gradien dari Z_Q kembali ke Z_E. Dengan cara ini kita tidak benar -benar meminimalkan fungsi kerugian tetapi masih dapat meneruskan beberapa informasi untuk pelatihan.

Fungsi kerugian

VQ-VAE menggunakan 3 kerugian untuk menghitung total kerugian selama pelatihan:

1. Kehilangan Rekonstruksi: Mengoptimalkan Decoder dan Encoder sebagai VAE, yaitu perbedaan antara gambar input dan rekonstruksi:

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

2. Kehilangan Codebook: Karena fakta bahwa gradien memotong embedding, algoritma pembelajaran kamus yang menggunakan kesalahan L2 untuk memindahkan vektor embedding E_I ke arah output enkoder digunakan.

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

   (SG mewakili operator gradien stop yang berarti tidak ada gradien mengalir melalui apa pun yang diterapkan)

3. Kehilangan komitmen: Karena volume ruang embedding tidak berdimensi, ia dapat tumbuh secara sewenang -wenang jika embeddings e_i tidak berlatih secepat parameter encoder, dan dengan demikian kehilangan komitmen ditambahkan untuk memastikan bahwa encoder berkomitmen untuk penyematan.

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

   (β adalah hiperparameter yang mengontrol seberapa besar kami ingin menimbang kehilangan komitmen dibandingkan dengan komponen lain)

Isi

1. Petunjuk Pengaturan
2. Melatih model Anda dari awal
3. Menghasilkan gambar dari model
4. Tinjauan Repositori
5. Hasil
   1. Gambar pelatihan
   2. Grafik pelatihan
   3. Menguji grafik
   4. Gambar yang dihasilkan
6. Pengamatan
7. Kredit

1. Instruksi Pengaturan

Anda dapat mengunduh repo atau mengkloningnya dengan menjalankan yang berikut di cmd prompt

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

2. Melatih model Anda dari awal

Anda dapat melatih model dari awal dengan perintah berikut (di 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 - Nama folder data
• data-folder - Nama folder data
• device - Atur perangkat (CPU atau CUDA, default: CPU)
• hidden-size - ukuran vektor laten (default: 40)
• k - Jumlah vektor laten (default: 512)
• batch-size - Ukuran Batch (Default: 128)
• num-epochs - Jumlah zaman (default: 10)
• lr - Tingkat Pembelajaran untuk Adam Optimizer (Default: 2E -4)
• beta - Kontribusi kehilangan komitmen, antara 0,1 dan 2.0 (default: 1.0)
• num-workers - Jumlah Pekerja untuk Lintasan Sampel (Default: CPU_Count () - 1)

Program secara otomatis mengunduh dataset MNIST dan menyimpannya di folder PATH_TO_MNIST_dataset (Anda perlu membuat folder ini). Ini hanya terjadi sekali.

Ini juga membuat folder logs dan folder models dan di dalamnya membuat folder dengan nama yang dilewatkan oleh Anda untuk menyimpan log dan model pos pemeriksaan di dalamnya masing -masing.

3. Menghasilkan gambar dari model

Untuk menghasilkan gambar baru dari Z sampel secara acak dari unit Gaussian menjalankan perintah berikut (di 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 - Nama file yang berisi model
• input - Mnist atau Random
• device - Atur perangkat (CPU atau CUDA, default: CPU)
• hidden-size - ukuran vektor laten (default: 40)
• k - Jumlah vektor laten (default: 512)
• filename - nama dengan file mana yang akan disimpan

Ini menghasilkan grid gambar 10*10 yang disimpan dalam folder bernama generatedImages .

Anda dapat menggunakan model pra-terlatih dengan mengunduhnya dari tautan di model.txt .

4. Tinjauan Repositori

Repositori berisi file berikut

• modules.py - berisi berbagai modul yang digunakan untuk membuat model kami
• VQ-VAE.py -berisi fungsi dan kode untuk melatih model VQ-VAE kami
• vector_quantizer.py - Kelas kuantisasi vektor didefinisikan dalam file ini
• generate-py -Menghasilkan gambar baru dari model pra-terlatih
• model.txt - Berisi tautan ke model pra -terlatih
• README.md - Readme memberikan gambaran umum tentang repo
• references.txt - referensi yang digunakan saat membuat repo ini
• readme_images - memiliki berbagai gambar untuk readme
• MNIST - Berisi dataset mnist zip (meskipun akan diunduh secara otomatis jika diperlukan)
• Training track for VQ-VAE.txt -berisi nilai kerugian selama pelatihan model VQ-VAE kami
• logs_VQ-VAE -berisi log tensorboard zip untuk model VQ-VAE kami (secara otomatis dibuat oleh program)
• testers.py - berisi beberapa fungsi untuk menguji modul yang kami tentukan

Perintah untuk menjalankan Tensorboard (di Google Colab):

 %load_ext tensorboard
%tensordboard --logdir [path_to_folder_with_logs]

5. Hasil

1. Gambar pelatihan

Gambar pelatihan

Gambar dari zaman ke -0

Gambar dari zaman ke -2

Gambar dari zaman ke -4

Gambar dari zaman ke -6

Gambar dari zaman ke -8

Gambar dari zaman ke -10

Rekonstruksi terus meningkat dan pada akhirnya hampir menyerupai gambar pelatihan_set yang tercermin dalam nilai kerugian (periksa Training track for VQ-VAE.txt ).

2. Grafik pelatihan

Kerugian rekonstruksi

Kerugian kuantisasi

Total_loss

Kehilangan total, kerugian rekonstruksi dan kerugian kuantisasi menurun secara seragam seperti yang diharapkan.

3. Grafik pengujian

Testing_loss

Kehilangan pengujian berkurang secara seragam seperti yang diharapkan.

4. Gambar yang dihasilkan

Kisi gambar berikut dihasilkan setelah melewati gambar mnist sebagai input:

Generasi ini cukup bagus.

Kisi -kisi gambar berikut dihasilkan setelah melewati AZ sampel secara acak dari unit Gaussian sebagai input ke model dan kemudian melewati decoder

Gambar tidak terlihat sempurna. Menyetel dimensi ruang laten, jumlah vektor embedding dll dapat membantu dalam menghasilkan gambar acak yang lebih baik.

6. Pengamatan

Model ini dilatih di Google Colab selama 10 zaman, dengan ukuran batch 128.

Setelah pelatihan, model ini dapat merekonstruksi gambar input dengan cukup baik, dan juga dapat menghasilkan gambar baru meskipun gambar yang dihasilkan tidak begitu baik.
Pelatihan serta kehilangan pengujian juga terus menurun hampir secara monoton.

Saya mengamati bahwa melatih model untuk lebih dari 10-20 zaman menghasilkan hasil yang menyarankan kemungkinan tanda overfitting dalam model. Juga, saya bereksperimen dengan dimensi yang berbeda dari ruang berlarut dan pada dimension = 40 menghasilkan hasil terbaik. Kisaran terbaik untuk dimensi keluar antara 16-42.

7. Kredit

Sumber -sumber berikut banyak membantu membuat repositori ini

• Pembelajaran Representasi Diskrit Neural - Aaron Van Den Oord, Oriol Vinyals, Koray Kavukcuoglu
• Menghasilkan beragam gambar kesetiaan tinggi dengan 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