<jupyter_start><jupyter_text>Model definition<jupyter_code>model = Model(latent_size=cfg.latent_size, data_size=cfg.data_size)
if cfg.variational == "flow":
    variational = VariationalFlow(
        latent_size=cfg.latent_size,
        data_size=cfg.data_size,
        flow_depth=cfg.flow_depth,
    )
elif cfg.variational == "mean-field":
    variational = VariationalMeanField(
        latent_size=cfg.latent_size, data_size=cfg.data_size
    )
else:
    raise ValueError(
        "Variational distribution not implemented: %s" % cfg.variational
    )
    
model.to(device)
variational.to(device)

optimizer = torch.optim.RMSprop(
    list(model.parameters()) + list(variational.parameters()),
    lr=cfg.learning_rate,
    centered=True,
)<jupyter_output><empty_output><jupyter_text>Dataset<jupyter_code>kwargs = {"num_workers": 4, "pin_memory": True} if cfg.use_gpu else {}
train_data, valid_data, test_data = load_binary_mnist(cfg, **kwargs)
best_valid_elbo = -np.inf
num_no_improvement = 0<jupyter_output>Downloading binary MNIST data...
Saved binary MNIST data to: /home/dm/projects/pytorch/variational-autoencoder/data/binary_mnist.h5<jupyter_text>Training<jupyter_code>for step, batch in enumerate(cycle(train_data)):
    x = batch[0].to(device)
    model.zero_grad()
    variational.zero_grad()
    z, log_q_z = variational(x, n_samples=1)
    log_p_x_and_z = model(z, x)
    # average over sample dimension
    elbo = (log_p_x_and_z - log_q_z).mean(1)
    # sum over batch dimension
    loss = -elbo.sum(0)
    loss.backward()
    optimizer.step()

    if step % cfg.log_interval == 0:
        print(
            f"step:\t{step}\ttrain elbo: {elbo.detach().cpu().numpy().mean():.2f}"
        )
        with torch.no_grad():
            valid_elbo, valid_log_p_x = evaluate(
                cfg.n_samples, model, variational, valid_data
            )
        print(
            f"step:\t{step}\t\tvalid elbo: {valid_elbo:.2f}\tvalid log p(x): {valid_log_p_x:.2f}"
        )
        if valid_elbo > best_valid_elbo:
            num_no_improvement = 0
            best_valid_elbo = valid_elbo
            states = {
                "model": model.state_dict(),
                "variational": variational.state_dict(),
            }
            torch.save(states, cfg.train_dir / "best_state_dict")
        else:
            num_no_improvement += 1

        if num_no_improvement > cfg.early_stopping_interval:
            checkpoint = torch.load(cfg.train_dir / "best_state_dict")
            model.load_state_dict(checkpoint["model"])
            variational.load_state_dict(checkpoint["variational"])
            with torch.no_grad():
                test_elbo, test_log_p_x = evaluate(
                    cfg.n_samples, model, variational, test_data
                )
            print(
                f"step:\t{step}\t\ttest elbo: {test_elbo:.2f}\ttest log p(x): {test_log_p_x:.2f}"
            )
            break<jupyter_output>step:	0	train elbo: -558.28
{'fraction_non_alphanumeric': 0.11709286675639301, 'fraction_numerical': 0.0070659488559892325, 'mean_word_length': 2.7069825436408976, 'pattern_counts': {'":': 6, '<': 11, '<?xml version=': 0, '>': 13, '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': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '195319', 'license': 'MIT', 'n_tokens_mistral': 1044, 'n_tokens_neox': 1011, 'n_words': 163, 'path': 'vae.ipynb', 'repo_name': 'dmhd1/variational-autoencoder', 'size': '2972'}