<jupyter_start><jupyter_text>Train a JointVAE model  Get MNIST dataBuild a simple JointVAE model on MNIST<jupyter_code>from utils.dataloaders import get_mnist_dataloaders
train_loader, test_loader = get_mnist_dataloaders(batch_size=64)<jupyter_output><empty_output><jupyter_text>Define latent distribution of the model<jupyter_code># Latent distribution will be joint distribution of 10 gaussian normal distributions
# and one 10 dimensional Gumbel Softmax distribution
latent_spec = {'cont': 10,
               'disc': [10]}<jupyter_output><empty_output><jupyter_text>Build a model<jupyter_code>from jointvae.models import VAE

model = VAE(latent_spec=latent_spec, img_size=(1, 32, 32))
print(model)<jupyter_output>VAE(
  (img_to_features): Sequential(
    (0): Conv2d(1, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
    (1): ReLU()
    (2): Conv2d(32, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
    (3): ReLU()
    (4): Conv2d(64, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
    (5): ReLU()
  )
  (features_to_hidden): Sequential(
    (0): Linear(in_features=1024, out_features=256, bias=True)
    (1): ReLU()
  )
  (fc_mean): Linear(in_features=256, out_features=10, bias=True)
  (fc_log_var): Linear(in_features=256, out_features=10, bias=True)
  (fc_alphas): ModuleList(
    (0): Linear(in_features=256, out_features=10, bias=True)
  )
  (latent_to_features): Sequential(
    (0): Linear(in_features=20, out_features=256, bias=True)
    (1): ReLU()
    (2): Linear(in_features=256, out_features=1024, bias=True)
    (3): ReLU()
  )
  (features_to_img): Sequential(
    (0): ConvTranspose2d(64, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
    (1): ReLU()
    (2): Conv[...]<jupyter_text>Train the model<jupyter_code>from torch import optim

# Build optimizer
optimizer = optim.Adam(model.parameters(), lr=5e-4)
from jointvae.training import Trainer

# Define the capacities
# Continuous channels
cont_capacity = [0.0, 5.0, 25000, 30.0]  # Starting at a capacity of 0.0, increase this to 5.0
                                         # over 25000 iterations with a gamma of 30.0
# Discrete channels
disc_capacity = [0.0, 5.0, 25000, 30.0]  # Starting at a capacity of 0.0, increase this to 5.0
                                         # over 25000 iterations with a gamma of 30.0

# Build a trainer
trainer = Trainer(model, optimizer,
                  cont_capacity=cont_capacity,
                  disc_capacity=disc_capacity)
# Build a visualizer which will be passed to trainer to visualize progress during training
from viz.visualize import Visualizer

viz = Visualizer(model)
# Train model for 10 epochs
# Note this should really be a 100 epochs and trained on a GPU, but this is just to demo

trainer.train(train_loader, epochs=10, save_training_gif=('./training.gif', viz))<jupyter_output>0/60000	Loss: 699.367
3200/60000	Loss: 526.442
6400/60000	Loss: 305.341
9600/60000	Loss: 279.369
12800/60000	Loss: 275.504
16000/60000	Loss: 273.062
19200/60000	Loss: 271.448
22400/60000	Loss: 271.727
25600/60000	Loss: 271.979
28800/60000	Loss: 268.752
32000/60000	Loss: 267.965
35200/60000	Loss: 268.570
38400/60000	Loss: 266.873
41600/60000	Loss: 266.212
44800/60000	Loss: 266.496
48000/60000	Loss: 265.427
51200/60000	Loss: 261.710
54400/60000	Loss: 262.949
57600/60000	Loss: 260.450
Epoch: 1 Average loss: 284.58
0/60000	Loss: 271.084
3200/60000	Loss: 259.601
6400/60000	Loss: 260.003
9600/60000	Loss: 257.882
12800/60000	Loss: 256.378
16000/60000	Loss: 259.008
19200/60000	Loss: 255.445
22400/60000	Loss: 256.978
25600/60000	Loss: 254.583
28800/60000	Loss: 256.848
32000/60000	Loss: 255.628
35200/60000	Loss: 251.737
38400/60000	Loss: 252.917
41600/60000	Loss: 253.391
44800/60000	Loss: 251.204
48000/60000	Loss: 248.861
51200/60000	Loss: 249.632
54400/60000	Loss: 250.133
57600/60000	Loss: 250.[...]<jupyter_text>Visualize<jupyter_code># Plot reconstructions
%matplotlib inline
import matplotlib.pyplot as plt

# Get a batch of data
for batch, labels in test_loader:
    break

# Reconstruct data using Joint-VAE model
recon = viz.reconstructions(batch)

plt.imshow(recon.numpy()[0, :, :], cmap='gray')
# Plot samples
samples = viz.samples()
plt.imshow(samples.numpy()[0, :, :], cmap='gray')
# Plot all traversals
traversals = viz.all_latent_traversals(size=10)
plt.imshow(traversals.numpy()[0, :, :], cmap='gray')
# Plot a grid of some traversals
traversals = viz.latent_traversal_grid(cont_idx=2, cont_axis=1, disc_idx=0, disc_axis=0, size=(10, 10))
plt.imshow(traversals.numpy()[0, :, :], cmap='gray')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1278638446432468, 'fraction_numerical': 0.16932140519310496, 'mean_word_length': 4.470167064439141, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 19, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '271827', 'license': 'MIT', 'n_tokens_mistral': 2320, 'n_tokens_neox': 1724, 'n_words': 441, 'path': 'train_model.ipynb', 'repo_name': 'RedHeadM/joint-vae', 'size': '4583'}