compact routine - separate to finetune

a mini compact routine for anyone who wants to compact their existing model. This is built into the new version of finetune anyway, but is here just to capture those outlier instances.
This commit is contained in:
erew123 authored and GitHub committed 2023-12-28 18:34:49 +00:00
1 parent 1a2b8571d0
commit 8b0ade5b9f
1 file changed
+26
@@ -0,0 +1,26 @@
import torch
import os
MODEL_PATH = "best_model.pth"
# Check if best_model.pth exists
if not os.path.exists(MODEL_PATH):
print(f"\nCannot find \033[94m{MODEL_PATH}\033[0m in the current folder.")
# Check if model.pth exists as an alternative
ALTERNATIVE_PATH = "model.pth"
if os.path.exists(ALTERNATIVE_PATH):
print(f"\nIf your model is currently called \033[93m{ALTERNATIVE_PATH}\033[0m please rename it to \033[94mbest_model.pth\033[0m.")
else:
print(f"Please put \033[91mbest_model.pth\033[0m in the current folder or rename your model file to \033[94mbest_model.pth\033[0m.")
exit(1)
# Attempt to load the model
try:
checkpoint = torch.load(MODEL_PATH, map_location=torch.device("cpu"))
except Exception as e:
print("Error loading checkpoint:", e)
raise
del checkpoint["optimizer"]
for key in list(checkpoint["model"].keys()):
if "dvae" in key:
del checkpoint["model"][key]
torch.save(checkpoint, "model.pth")
print(f"\nYour compacted model is created \033[93mmodel.pth\033[0m.")