Load a LitSegger model from a checkpoint.
Parameters
checkpoint_path : str
Specific checkpoint file to load, or directory where the model checkpoints are stored.
If directory, the latest checkpoint is loaded.
Returns
LitSegger
The loaded LitSegger model.
Raises
FileNotFoundError
If the specified checkpoint file does not exist.
Source code in src/segger/prediction/predict.py
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96 | def load_model(checkpoint_path: str) -> LitSegger:
"""
Load a LitSegger model from a checkpoint.
Parameters
----------
checkpoint_path : str
Specific checkpoint file to load, or directory where the model checkpoints are stored.
If directory, the latest checkpoint is loaded.
Returns
-------
LitSegger
The loaded LitSegger model.
Raises
------
FileNotFoundError
If the specified checkpoint file does not exist.
"""
checkpoint_path = Path(checkpoint_path)
msg = f"No checkpoint found at {checkpoint_path}. Please make sure you've provided the correct path."
# Get last checkpoint if directory is provided
if os.path.isdir(checkpoint_path):
checkpoints = glob.glob(str(checkpoint_path / "*.ckpt"))
if len(checkpoints) == 0:
raise FileNotFoundError(msg)
# Sort checkpoints by epoch and step
def sort_order(c):
match = re.match(r".*epoch=(\d+)-step=(\d+).ckpt", c)
return int(match[1]), int(match[2])
checkpoint_path = Path(sorted(checkpoints, key=sort_order)[-1])
elif not checkpoint_path.exists():
raise FileExistsError(msg)
# Load model from checkpoint
lit_segger = LitSegger.load_from_checkpoint(
checkpoint_path=checkpoint_path,
)
return lit_segger
|