<jupyter_start><jupyter_text>Variational Autoecoder<jupyter_code>%matplotlib inline
import json

import numpy as np
import matplotlib.pyplot as plt
from torch.optim import Adam
from torch.utils.data.dataloader import DataLoader
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor

from callbacks import PlotCallback
from loss import loss_criterion
from model import VAE
from train import Trainer
from utils import load_model, save_model
from config import *<jupyter_output><empty_output><jupyter_text>Let's set training and model parameters<jupyter_code># Training params
batch_size = 64
lr = 0.001
device = "cuda"

# Model params
input_size = 784
hidden_size = 512
latent_size = 10<jupyter_output><empty_output><jupyter_text>Load data set and create data loader<jupyter_code># Load data set
dataset = MNIST(DATA_DIR, transform=ToTensor(), download=True)

# Create data loader
data_loader = DataLoader(dataset=dataset, batch_size=batch_size)<jupyter_output><empty_output><jupyter_text>Initialize model and optimizer<jupyter_code># Initialize model
model = VAE(input_size, hidden_size, latent_size)

# Initialize optimizer
optimizer = Adam(model.parameters(), lr=lr)<jupyter_output><empty_output><jupyter_text>Now that we have our dataset, model and optimizer ready, only that's left is to initialize it in `Trainer` and run training.<jupyter_code># Initialize trainer
trainer = Trainer(model=model,
                  data_loader=data_loader,
                  optimizer=optimizer,
                  device=device)<jupyter_output><empty_output><jupyter_text>Run training for number of epochs.<jupyter_code>losses = trainer.run_train_loop(epochs=50)<jupyter_output><empty_output><jupyter_text>Plot losses during training<jupyter_code>plt.plot(np.arange(len(losses)), losses)
plt.title("Training")
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.show()<jupyter_output><empty_output><jupyter_text>Even though it's used as a callback to plot samples during training, it can be used to generate samples aside from the train loop.<jupyter_code>plot_cb = PlotCallback()
for _ in range(5):
    plot_cb(trainer)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.08387395736793327, 'fraction_numerical': 0.007877664504170528, 'mean_word_length': 6.794223826714801, 'pattern_counts': {'":': 0, '<': 33, '<?xml version=': 0, '>': 33, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '662615', 'license': 'Apache-2.0', 'n_tokens_mistral': 716, 'n_tokens_neox': 684, 'n_words': 178, 'path': 'notebooks/train_and_eval.ipynb', 'repo_name': 'bvezilic/Variational-autoencoder', 'size': '2158'}