append sims in batches

[enricobarausse]

Hi,

I have a very heavy set of simulations which causes my GPU to go OOM when I append them. I tried appending them in batches, do the training on one batch, load another batch, train again etc. The code is below (device is cuda). the problem is that I still OOM, because it seems that inference.append_simulations saves all previous simulations, so RAM usage accumulates. Is it possible to purge the old simulations while saving the network weights? also, I set trained_once=True after the first batch training because the second batch is still sampled from the prior, is that correct?

theta_t = torch.as_tensor(theta, dtype=torch.float32, device=device)
x_t     = torch.as_tensor(HH,    dtype=torch.float32, device=device)

prior = BoxUniform(low=PRIOR_LOW_T.to(device), high=PRIOR_HIGH_T.to(device), device=device)

inference = NPSE(prior, sde_type="ve",device=device)


N = theta.shape[0]
trained_once=False

for start in range(0, N, CHUNK):
    sl = slice(start, min(start + CHUNK, N))

    theta_chunk = torch.as_tensor(theta[sl], dtype=torch.float32, device=device)
    x_chunk     = torch.as_tensor(HH[sl],    dtype=torch.float32, device=device)

    # append this batch of sims
    inference.append_simulations(theta_chunk, x_chunk)

    # train incrementally
    _ = inference.train(
        max_num_epochs=max_epochs,
        training_batch_size=batch_size,
        force_first_round_loss=trained_once,
        stop_after_epochs=20
        )

    trained_once=True

    print(f"[OK] Trained on chunk {start}:{sl.stop}")

[michaeldeistler]

Hey! Thanks for reaching out. A quick fix might be to do:

inference.append_simulations(theta_chunk, x_chunk, data_device="cpu")

By doing this, data gets stored on CPU, and only individual batches are copied to GPU.

If this does not help, sbi also allows you to write your own training loop from scratch, with full flexibility for, e.g., dataloading. See here for a tutorial.

Hope that helps!
Michael