peal.architectures.predictors.get_predictor

peal.architectures.predictors.get_predictor(predictor, device='cuda')[source]

Resolve a predictor reference into a model on device.

Accepts the different ways a predictor is referred to throughout PEAL:

  • an nn.Module or any callable, which is returned unchanged;

  • a path ending in .cpl, loaded with torch.load (retried with weights_only=False for full-pickle checkpoints);

  • a path ending in .onnx, converted to a trainable torch module by peal.architectures.onnx_predictor.load_onnx_predictor;

  • anything else, which is read as a predictor config. The config either names weights_path (a TorchvisionModel of config.architecture with config.task.output_channels outputs is built and the state dict loaded into it), selects the pretrained torchvision_resnet18_imagenet, or points at a run directory whose model.cpl is loaded.

Parameters:

  • predictor (nn.Module or callable or str or PredictorConfig) – The reference to resolve.

  • device (str, optional) – Device the model is moved to.

Returns:

  • model (nn.Module or callable) – The resolved predictor.

  • config (PredictorConfig or None) – The config it was built from, or None when the predictor was given directly or as a checkpoint path.