JimmyChin1998/Pytorch-Learning-File

1"""
2Contains various utility functions for PyTorch model training and saving.
3"""
4import torch
5from pathlib import Path
6
7def save_model(model: torch.nn.Module,
8               target_dir: str,
9               model_name: str):
10  """Saves a PyTorch model to a target directory.
11
12  Args:
13    model: A target PyTorch model to save.
14    target_dir: A directory for saving the model to.
15    model_name: A filename for the saved model. Should include
16      either ".pth" or ".pt" as the file extension.
17
18  Example usage:
19    save_model(model=model_0,
20               target_dir="models",
21               model_name="05_going_modular_tingvgg_model.pth")
22  """
23  # Create target directory
24  target_dir_path = Path(target_dir)
25  target_dir_path.mkdir(parents=True,
26                        exist_ok=True)
27
28  # Create model save path
29  assert model_name.endswith(".pth") or model_name.endswith(".pt"), "model_name should end with '.pt' or '.pth'"
30  model_save_path = target_dir_path / model_name
31
32  # Save the model state_dict()
33  print(f"[INFO] Saving model to: {model_save_path}")
34  torch.save(obj=model.state_dict(),
35             f=model_save_path)
36