garmentiq.classification.load_model

Loading a trained classification model.

View Source

 1"""Loading a trained classification model."""
 2import torch
 3import torch.nn as nn
 4import torch.nn.functional as F
 5from typing import Type, List, Union
 6from garmentiq.utils.device import resolve_device
 7from garmentiq.utils.checkpoint import load_state_dict_checked
 8
 9
10def load_model(
11    model_path: str,
12    model_class: Type[nn.Module],
13    model_args: dict,
14    device: Union[str, torch.device] = "cpu",
15):
16    """
17    Loads a PyTorch model from a checkpoint and prepares it for inference.
18
19    This function initializes a model from the provided `model_class`, loads its weights from
20    the given file path, moves it to the requested device, and sets it to evaluation mode.
21
22    Args:
23        model_path (str): Path to the saved model checkpoint (.pth or .pt file).
24        model_class (Type[nn.Module]): The class definition of the model to be instantiated.
25                                       This must be a subclass of `torch.nn.Module`.
26        model_args (dict): A dictionary of arguments used to initialize the model class.
27        device (Union[str, torch.device], optional): The device to load the model onto, e.g.
28                                                     `"cpu"`, `"cuda"`, `"cuda:0"`, or `"mps"`.
29                                                     Hardware acceleration is opt-in; pass it
30                                                     explicitly to use a GPU or Apple Silicon.
31                                                     Default is `"cpu"`.
32
33    Raises:
34        ValueError: If the requested `device` is invalid or unavailable on this machine.
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(**model_args).to(device)
42    state_dict = torch.load(model_path, map_location=device, weights_only=True)
43    new_state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
44    load_state_dict_checked(model, new_state_dict, model_path)
45    model.eval()
46
47    return model

def load_model( model_path: str, model_class: Type[torch.nn.modules.module.Module], model_args: dict, device: Union[str, torch.device] = 'cpu'): View Source

11def load_model(
12    model_path: str,
13    model_class: Type[nn.Module],
14    model_args: dict,
15    device: Union[str, torch.device] = "cpu",
16):
17    """
18    Loads a PyTorch model from a checkpoint and prepares it for inference.
19
20    This function initializes a model from the provided `model_class`, loads its weights from
21    the given file path, moves it to the requested device, and sets it to evaluation mode.
22
23    Args:
24        model_path (str): Path to the saved model checkpoint (.pth or .pt file).
25        model_class (Type[nn.Module]): The class definition of the model to be instantiated.
26                                       This must be a subclass of `torch.nn.Module`.
27        model_args (dict): A dictionary of arguments used to initialize the model class.
28        device (Union[str, torch.device], optional): The device to load the model onto, e.g.
29                                                     `"cpu"`, `"cuda"`, `"cuda:0"`, or `"mps"`.
30                                                     Hardware acceleration is opt-in; pass it
31                                                     explicitly to use a GPU or Apple Silicon.
32                                                     Default is `"cpu"`.
33
34    Raises:
35        ValueError: If the requested `device` is invalid or unavailable on this machine.
36
37    Returns:
38        torch.nn.Module: The loaded and ready-to-use model, placed on `device`.
39    """
40    device = resolve_device(device)
41
42    model = model_class(**model_args).to(device)
43    state_dict = torch.load(model_path, map_location=device, weights_only=True)
44    new_state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
45    load_state_dict_checked(model, new_state_dict, model_path)
46    model.eval()
47
48    return model

Loads a PyTorch model from a checkpoint and prepares 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.

Arguments:

model_path (str): Path to the saved model checkpoint (.pth or .pt file).
model_class (Type[nn.Module]): The class definition of the model to be instantiated. This must be a subclass of torch.nn.Module.
model_args (dict): A dictionary of arguments used to initialize the model class.
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.

Returns:

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