Spaces:
Running on Zero
Running on Zero
File size: 2,025 Bytes
0122a25 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 | """Convert model weights for release."""
from __future__ import annotations
import argparse
import hashlib
import os
import torch
def save_weights_with_hash(
state_dict: dict[str, torch.Tensor],
path: str,
filename: str,
digits: int = 6,
) -> None:
"""Saves the model weights and append a 6-digit hash to the filename.
Args:
state_dict (dict[str, torch.Tensor]): The model weights to save.
path (str): The directory path to save the model.
filename (str): The filename to save the model.
digits (int, optional): The number of digits to use for the hash.
Defaults to 6.
"""
os.makedirs(path, exist_ok=True)
with open(os.path.join(path, filename), "wb") as f:
torch.save(state_dict, f)
# Create a hash of the file
sha256_hash = hashlib.sha256()
with open(os.path.join(path, filename), "rb") as f:
for byte_block in iter(lambda: f.read(4096), b""):
sha256_hash.update(byte_block)
# Get the hexadecimal representation of the hash
short_hash = sha256_hash.hexdigest()[:digits]
os.rename(
os.path.join(path, filename),
os.path.join(path, f"{filename}_{short_hash}.pt"),
)
def main() -> None:
"""Main function."""
parser = argparse.ArgumentParser(
description="Save trained model checkpoint with a filename hash."
)
parser.add_argument("path", type=str, help="The path to the checkpoint.")
parser.add_argument(
"--outdir",
type=str,
help="The path to output the model.",
default="./work_dir/release",
)
parser.add_argument(
"--name", type=str, help="The base name of the released file."
)
args = parser.parse_args()
checkpoint = torch.load(
args.path, weights_only=False, map_location=torch.device("cpu")
)
state_dict = {"state_dict": checkpoint["state_dict"]}
save_weights_with_hash(state_dict, args.outdir, args.name)
if __name__ == "__main__":
main()
|