Utilities

forward_in_batches(model, inputs, device, compute_grad=False, batch_size=5000, use_tqdm=False)

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	ndarray | Tensor	Input array to be fed to the model	required
device	str	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 5_000.	5000
use_tqdm	bool	Whether to display a tqdm progress bar during computation. Defaults to False.	False

Returns:

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

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.

gradient(inp_tensor, out_tensor)

Computes the gradient of the output tensor with respect to the input tensor using pytorch's autograd.

Parameters:

Name	Type	Description	Default
inp_tensor	Tensor	input tensor	required
out_tensor	Tensor	output tensor	required

Returns:

Type	Description
Tensor	torch.Tensor: the gradient