60 lines
1.9 KiB
Python
60 lines
1.9 KiB
Python
import os
|
|
import torch
|
|
import json
|
|
from pathlib import Path
|
|
import safetensors
|
|
import wandb
|
|
|
|
|
|
def create_folder_if_necessary(path):
|
|
path = "/".join(path.split("/")[:-1])
|
|
Path(path).mkdir(parents=True, exist_ok=True)
|
|
|
|
|
|
def safe_save(ckpt, path):
|
|
try:
|
|
os.remove(f"{path}.bak")
|
|
except OSError:
|
|
pass
|
|
try:
|
|
os.rename(path, f"{path}.bak")
|
|
except OSError:
|
|
pass
|
|
if path.endswith(".pt") or path.endswith(".ckpt"):
|
|
torch.save(ckpt, path)
|
|
elif path.endswith(".json"):
|
|
with open(path, "w", encoding="utf-8") as f:
|
|
json.dump(ckpt, f, indent=4)
|
|
elif path.endswith(".safetensors"):
|
|
safetensors.torch.save_file(ckpt, path)
|
|
else:
|
|
raise ValueError(f"File extension not supported: {path}")
|
|
|
|
|
|
def load_or_fail(path, wandb_run_id=None):
|
|
accepted_extensions = [".pt", ".ckpt", ".json", ".safetensors"]
|
|
try:
|
|
assert any(
|
|
[path.endswith(ext) for ext in accepted_extensions]
|
|
), f"Automatic loading not supported for this extension: {path}"
|
|
if not os.path.exists(path):
|
|
checkpoint = None
|
|
elif path.endswith(".pt") or path.endswith(".ckpt"):
|
|
checkpoint = torch.load(path, map_location="cpu")
|
|
elif path.endswith(".json"):
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
checkpoint = json.load(f)
|
|
elif path.endswith(".safetensors"):
|
|
checkpoint = {}
|
|
with safetensors.safe_open(path, framework="pt", device="cpu") as f:
|
|
for key in f.keys():
|
|
checkpoint[key] = f.get_tensor(key)
|
|
return checkpoint
|
|
except Exception as e:
|
|
if wandb_run_id is not None:
|
|
wandb.alert(
|
|
title=f"Corrupt checkpoint for run {wandb_run_id}",
|
|
text=f"Training {wandb_run_id} tried to load checkpoint {path} and failed",
|
|
)
|
|
raise e
|