garmentiq.landmark.detection.load_model

Loading the HRNet landmark detection model onto a device.

 1"""Loading the HRNet landmark detection model onto a device."""
 2import torch
 3from typing import Callable, Type, Union
 4from garmentiq.utils.device import resolve_device
 5from garmentiq.utils.checkpoint import load_state_dict_checked
 6
 7
 8def load_model(
 9    model_path: str,
10    model_class: Type[torch.nn.Module],
11    device: Union[str, torch.device] = "cpu",
12):
13    """
14    Load a PyTorch model from a checkpoint and prepare it for inference.
15
16    This function initializes a model from the provided `model_class`, loads its weights from
17    the given file path, moves it to the requested device, and sets it to evaluation mode.
18    When multiple CUDA GPUs are available and a CUDA device is requested, the model is
19    additionally wrapped with `DataParallel` for multi-GPU inference.
20
21    Args:
22        model_path (str): Path to the saved model checkpoint (.pth or .pt file).
23        model_class (Type[torch.nn.Module]): The class definition of the model to be instantiated.
24                                           This must be a subclass of `torch.nn.Module`.
25        device (Union[str, torch.device], optional): The device to load the model onto, e.g.
26                                                     `"cpu"`, `"cuda"`, `"cuda:0"`, or `"mps"`.
27                                                     Hardware acceleration is opt-in; pass it
28                                                     explicitly to use a GPU or Apple Silicon.
29                                                     Default is `"cpu"`.
30
31    Raises:
32        ValueError: If the requested `device` is invalid or unavailable on this machine.
33        RuntimeError: If the model checkpoint cannot be loaded.
34
35    Returns:
36        torch.nn.Module: The loaded and ready-to-use model, placed on `device`.
37    """
38    device = resolve_device(device)
39
40    model = model_class
41    load_state_dict_checked(
42        model, torch.load(model_path, map_location=device), model_path
43    )
44    model = model.to(device)
45
46    if device.type == "cuda" and torch.cuda.device_count() > 1:
47        model = torch.nn.DataParallel(model)
48
49    model.eval()
50
51    return model
def load_model( model_path: str, model_class: Type[torch.nn.modules.module.Module], device: Union[str, torch.device] = 'cpu'): 
 9def load_model(
10    model_path: str,
11    model_class: Type[torch.nn.Module],
12    device: Union[str, torch.device] = "cpu",
13):
14    """
15    Load a PyTorch model from a checkpoint and prepare it for inference.
16
17    This function initializes a model from the provided `model_class`, loads its weights from
18    the given file path, moves it to the requested device, and sets it to evaluation mode.
19    When multiple CUDA GPUs are available and a CUDA device is requested, the model is
20    additionally wrapped with `DataParallel` for multi-GPU inference.
21
22    Args:
23        model_path (str): Path to the saved model checkpoint (.pth or .pt file).
24        model_class (Type[torch.nn.Module]): The class definition of the model to be instantiated.
25                                           This must be a subclass of `torch.nn.Module`.
26        device (Union[str, torch.device], optional): The device to load the model onto, e.g.
27                                                     `"cpu"`, `"cuda"`, `"cuda:0"`, or `"mps"`.
28                                                     Hardware acceleration is opt-in; pass it
29                                                     explicitly to use a GPU or Apple Silicon.
30                                                     Default is `"cpu"`.
31
32    Raises:
33        ValueError: If the requested `device` is invalid or unavailable on this machine.
34        RuntimeError: If the model checkpoint cannot be loaded.
35
36    Returns:
37        torch.nn.Module: The loaded and ready-to-use model, placed on `device`.
38    """
39    device = resolve_device(device)
40
41    model = model_class
42    load_state_dict_checked(
43        model, torch.load(model_path, map_location=device), model_path
44    )
45    model = model.to(device)
46
47    if device.type == "cuda" and torch.cuda.device_count() > 1:
48        model = torch.nn.DataParallel(model)
49
50    model.eval()
51
52    return model

Load a PyTorch model from a checkpoint and prepare it for inference.

This function initializes a model from the provided model_class, loads its weights from the given file path, moves it to the requested device, and sets it to evaluation mode. When multiple CUDA GPUs are available and a CUDA device is requested, the model is additionally wrapped with DataParallel for multi-GPU inference.

Arguments:
  • model_path (str): Path to the saved model checkpoint (.pth or .pt file).
  • model_class (Type[torch.nn.Module]): The class definition of the model to be instantiated. This must be a subclass of torch.nn.Module.
  • device (Union[str, torch.device], optional): The device to load the model onto, e.g. "cpu", "cuda", "cuda:0", or "mps". Hardware acceleration is opt-in; pass it explicitly to use a GPU or Apple Silicon. Default is "cpu".
Raises:
  • ValueError: If the requested device is invalid or unavailable on this machine.
  • RuntimeError: If the model checkpoint cannot be loaded.
Returns:

torch.nn.Module: The loaded and ready-to-use model, placed on device.