Multi-round inference and GPU memory usage

[ali-akhavan89]

Hello,

I have noticed that GPU memory usage increases each time a new round of training starts using SNPE-C, until it eventually reaches the maximum. I wondered if this is a normal and expected behavior. Could you please let me know what is causing the increase in memory usage in each round? Would you please share some insights into ways to avoid such memory problems?

Thank you,
Ali

[janfb]

Hallo @ali-akhavan89 ,

the reason is that each round new simulations are added to the inference object and saved internally here:

sbi/sbi/inference/trainers/npe/npe_base.py

Lines 191 to 193 in 9152e93

	self._theta_roundwise.append(theta)
	self._x_roundwise.append(x)
	self._prior_masks.append(prior_masks)

In each round, all data from all previous rounds is then used for training, see here:

sbi/sbi/inference/trainers/npe/npe_base.py

Lines 294 to 315 in 9152e93

	start_idx = int(discard_prior_samples and self._round > 0)
	# For non-atomic loss, we can not reuse samples from previous rounds as of now.
	# SNPE-A can, by construction of the algorithm, only use samples from the last
	# round. SNPE-A is the only algorithm that has an attribute `_ran_final_round`,
	# so this is how we check for whether or not we are using SNPE-A.
	if self.use_non_atomic_loss or hasattr(self, "_ran_final_round"):
	start_idx = self._round
	# Set the proposal to the last proposal that was passed by the user. For
	# atomic SNPE, it does not matter what the proposal is. For non-atomic
	# SNPE, we only use the latest data that was passed, i.e. the one from the
	# last proposal.
	proposal = self._proposal_roundwise[-1]
	train_loader, val_loader = self.get_dataloaders(
	start_idx,
	training_batch_size,
	validation_fraction,
	resume_training,
	dataloader_kwargs=dataloader_kwargs,
	)

There are ways to circumvent this, e.g., by resetting self._theta_roundwise and self._x_roundwise in each round, but you need to take care of the start_idx then as well.

A more convenient way for more control might be implementing your own training loop like described here:

https://github.com/sbi-dev/sbi/blob/main/tutorials/18_training_interface.ipynb

Does this help?

[ali-akhavan89]

Thanks a lot, Jan! It makes sense. I overlooked the impact of append_simulations. Also, thanks a lot for the suggestions. They are very helpful.