Feature extraction code puts models into eval() and inference_mode() behind the scenes without printing an explicit message

[GeorgeBatch]

• TIA Toolbox version: develop branch
• Python version: 3.11
• Operating System: linux

Description

I realised that in my code, like in the foundation models example notebook, we do not explicitly put the model into the .eval() mode or activate torch.no_grad() or torch.inference_mode() context managers.

The same is true for the older example notebooks on which I based my code:

• extract_deep_features() and extract_composition_features() in https://github.com/TissueImageAnalytics/tiatoolbox/blob/develop/examples/full-pipelines/slide-graph.ipynb
• extract_deep_features() and extract_composition_features() in https://github.com/TissueImageAnalytics/tiatoolbox/blob/develop/examples/inference-pipelines/slide-graph.ipynb

I got worried that I was using models in train mode and that I would need to recompute all the features I computed over the last couple of weeks, but then I checked the source code and realised that we rely on the _infer_batch() function, which runs model.eval() for all models that use it (line 173):

tiatoolbox/tiatoolbox/models/architecture/vanilla.py

Lines 148 to 178 in 4a1940d

	def _infer_batch(
	model: nn.Module,
	batch_data: torch.Tensor,
	device: str,
	) -> dict[str, np.ndarray]:
	"""Run inference on an input batch.
	Contains logic for forward operation as well as i/o aggregation.
	Args:
	model (nn.Module):
	PyTorch defined model.
	batch_data (torch.Tensor):
	A batch of data generated by
	`torch.utils.data.DataLoader`.
	device (str):
	Transfers model to the specified device. Default is "cpu".
	"""
	img_patches_device = batch_data.to(device=device).type(
	torch.float32,
	) # to NCHW
	img_patches_device = img_patches_device.permute(0, 3, 1, 2).contiguous()
	# Inference mode
	model.eval()
	# Do not compute the gradient (not training)
	with torch.inference_mode():
	output = model(img_patches_device)
	# Output should be a single tensor or scalar
	return output.cpu().numpy()

Suggestion

I think this behaviour is as it should be, but maybe adding a message specifying that the models are being put into eval mode would be beneficial.