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:
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.")
|
||||
Reference in new issue
Block a user