85 lines
1.4 KiB
Python
85 lines
1.4 KiB
Python
import os
|
|
import torch
|
|
|
|
|
|
def save_checkpoint(
|
|
model,
|
|
optimizer,
|
|
epoch,
|
|
loss,
|
|
path,
|
|
):
|
|
"""
|
|
Save a training checkpoint.
|
|
"""
|
|
|
|
os.makedirs(os.path.dirname(path), exist_ok=True)
|
|
|
|
torch.save(
|
|
{
|
|
"epoch": epoch,
|
|
"model_state_dict": model.state_dict(),
|
|
"optimizer_state_dict": optimizer.state_dict(),
|
|
"loss": loss,
|
|
},
|
|
path,
|
|
)
|
|
|
|
|
|
def load_checkpoint(
|
|
path,
|
|
model,
|
|
optimizer=None,
|
|
map_location="cpu",
|
|
):
|
|
"""
|
|
Load a training checkpoint.
|
|
|
|
Returns:
|
|
epoch (int)
|
|
loss (float)
|
|
"""
|
|
|
|
checkpoint = torch.load(
|
|
path,
|
|
map_location=map_location,
|
|
)
|
|
|
|
model.load_state_dict(
|
|
checkpoint["model_state_dict"]
|
|
)
|
|
|
|
if optimizer is not None:
|
|
optimizer.load_state_dict(
|
|
checkpoint["optimizer_state_dict"]
|
|
)
|
|
|
|
return (
|
|
checkpoint["epoch"],
|
|
checkpoint["loss"],
|
|
)
|
|
|
|
|
|
def latest_checkpoint(folder):
|
|
"""
|
|
Returns the newest checkpoint file in a folder.
|
|
"""
|
|
|
|
if not os.path.exists(folder):
|
|
return None
|
|
|
|
checkpoints = [
|
|
os.path.join(folder, file)
|
|
for file in os.listdir(folder)
|
|
if file.endswith(".pth")
|
|
]
|
|
|
|
if len(checkpoints) == 0:
|
|
return None
|
|
|
|
checkpoints.sort(
|
|
key=os.path.getmtime,
|
|
reverse=True,
|
|
)
|
|
|
|
return checkpoints[0] |