Utilities

forward_in_batches(model, inputs, device, compute_grad=False, batch_size=2 ** 16, use_tqdm=False, full_precision=True)

Runs the model model onto the inputs array and returns the computed output. As the input array may be large, the function splits it in batches that are fed one after the other to the model.

Parameters:

Name	Type	Description	Default
model	Module	the pytorch model to run	required
inputs	Tensor	Input array to be fed to the model. Also supports numpy array, which will be converted automatically.	required
device	device	which device the computation should be run on.	required
compute_grad	bool	whether to also compute and returns the gradients of the model at the input points. Defaults to False.	False
batch_size	int	Batch size for the forward computation. Defaults to 2**16.	2 ** 16
use_tqdm	bool	Whether to display a tqdm progress bar during computation. Defaults to False.	False
full_precision	bool	Whether to compute with torch.float32 (more accurate) or torch.float16 (faster but some slight errors may occur). Defaults to True.	True

Returns:

Type	Description
Tensor	torch.Tensor: the output model(inputs). If compute_grad is set to True, also returns the gradients as another tensor.

get_device(force_cpu=False)

Utility function for selecting the correct pytorch device. Returns the first gpu device found on the current system if it exists and "cpu" otherwise.

Parameters:

Name	Type	Description	Default
force_cpu	bool	Forces the computation to happen on the cpu. If set to False, the device returned will always be torch.device("cpu"). Defaults to False.	False

Returns:

Type	Description
device	torch.device: a torch device on which to run the program.