Training an embedding net in batches

[btetaren]

Hello,

I am trying to train an embedding net with a large data-set. I am using an embedding net because my data involves higher-dimensional images. As a result, the training data is very large (too large to be stored in memory, I am working with CPUs). However, I have not found a way to write a custom training loop when using an embedding network. Is there a way to write a custom training loop when using an embedding network? If this is not possible, how would I go about training the embedding net in batches? Any help would be much appreciated. Thank you.

[michaeldeistler]

Hi there,

please follow this tutorial. As you can see there, it fully exposes the dataloader, which you can modify to your needs. For the embedding net, do:

from sbi.neural_nets.net_builders import build_nsf

embedding_net = ...
density_estimator = build_nsf(theta, x, embedding_net=embedding_net)

See here for more neural net options.

Hope this helps!
Michael

[btetaren]

ok, below goes through the general process of what I am doing. Let me know if you would like the full trace back of the error as well.

import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.nn.functional as F

from sbi import analysis, utils
from sbi.inference import NPE, simulate_for_sbi
from sbi.utils.user_input_checks import (
    check_sbi_inputs,
    process_prior,
    process_simulator,
)

seed = 0
torch.manual_seed(seed);

from sbi.neural_nets.embedding_nets import CNNEmbedding
from sbi.neural_nets.net_builders import build_nsf
from sbi.inference.posteriors import MCMCPosterior
from sbi.inference.potentials import likelihood_estimator_based_potential

from torch.optim import Adam, AdamW
from torch.utils.data import Dataset

#Simulations are done seperately and saved as tensors, MyDataset loads theta,x 
# pairs from disc. Theta is a (10,10) image, flattened, of shape (100,). 


#Prior is a uniform distribution, event_shape is 100
img=torch.ones((10,10)).flatten()        
prior = utils.BoxUniform(low=3.5*img,high=5.5*img)  


embedding_net = CNNEmbedding(
    input_shape=(10,2000),
    in_channels=1,
    kernel_size=1,
    pool_kernel_size=1
)

# Create an instance of the dataset
train_data = MyDataset(data_dir=theta_train_path,data_dir2=x_train_path)

# Create a DataLoader
train_loader = torch.utils.data.DataLoader(train_data, batch_size=10000)

#create a second dataset and use only parts of the data to construct the density estimator
dummy_data = MyDataset(data_dir=theta_dummy_path,data_dir2=x_dummy_path)
dummy_loader = torch.utils.data.DataLoader(dummy_data, batch_size=1000)
dummy_theta, dummy_x = next(iter(dummy_loader))


# setup the inference procedure using dummy data set
density_estimator = build_nsf(dummy_theta,dummy_x,embedding_net=embedding_net)

#optimizer
optw = AdamW(list(density_estimator.parameters()), lr=5e-4)

#training loop
num_epochs = 300
for ep in range(num_epochs):
    for _, (theta_batch, x_batch) in enumerate(train_loader):
        optw.zero_grad()
        losses = density_estimator.loss(theta_batch, condition=x_batch)
        loss = torch.mean(losses)
        loss.backward()
        optw.step()
    if ep % 10 == 0:
        print("last loss", loss.item())


#x is a (10,2000) set of 1D spectra, flattened, of shape (20000,),
#loaded from test data on disc
x_observed=torch.randn(20000)


# Build the posterior.
potential, tf = likelihood_estimator_based_potential(density_estimator, prior, x_observed)
posterior = MCMCPosterior(
    potential,
    proposal=prior,
    theta_transform=tf,
    num_chains=1,
    thin=1,
    warmup_steps=100,
    method="nuts_pymc"
)

#sample
samples = posterior.sample((10000,), x=x_observed)```

[michaeldeistler]

Thanks! I am confused what you want to achieve now though. You are currently using NPE, because you are training a conditional density estimator to estimate the posterior p(theta | x):

density_estimator = build_nsf(dummy_theta,dummy_x,embedding_net=embedding_net)

However, you are then using likelihood_estimator_based_potential, which is meant to be combined with NLE (not NPE). This is why you are getting the error.

When using NPE, why do you even want to use the MCMCPosterior? Why not just sample with the DirectPosterior, or even just from the density_estimator?

Michael

[btetaren]

Oh ok, I was mixing up NPE and NLE, sorry.

If I wanted to sample:

a) directly from the density estimator: do I use density_estimator.sample((10000,),x=x_observed)? or do I have to use build_posterior first?

b) using DirectPosterior: posterior = DirectPosterior(density_estimator, prior=prior) and samples = posterior.sample((10000,), x=x_observed)?

Also, can I still set the sample_with option or does DirectPosterior just use rejection sampling? and in the case of directly from the density estimator, what type of sampling is used here?

thanks.

[michaeldeistler]

This samples directly from the trained normalizing flow:

density_estimator.sample((10000,), condition=x_observed)

The following also samples directly from the normalizing flow, but it then rejects samples that are outside of the support of the prior:

posterior = DirectPosterior(density_estimator, prior=prior)
samples = posterior.sample((10000,), x=x_observed)

Hope this helps!
Michael

[btetaren]

ok, I think I want the DirectPosterior then, thanks!