janitor

janitor
git clone https://git.prime8.dev/janitor.git
Log | Files | Refs | README | LICENSE

create.py (713B)

      1 from os import path
      2 from wget import download
      3 from torch import load, jit
      4 from model import ASTModel
      5 
      6 model = ASTModel()
      7 model.eval()
      8 
      9 checkpoint_path = "./pretrained_models/audioset_10_10_0.4593.pth"
     10 if not path.exists("./pretrained_models/audioset_10_10_0.4593.pth"):
     11     audioset_mdl_url = (
     12         "https://www.dropbox.com/s/cv4knew8mvbrnvq/audioset_0.4593.pth?dl=1"
     13     )
     14     download(
     15         audioset_mdl_url,
     16         out=checkpoint_path,
     17     )
     18 checkpoint = load(checkpoint_path, weights_only=True)
     19 state_dict = {
     20     key.replace("module.", ""): value for key, value in checkpoint.items()
     21 }
     22 model.load_state_dict(state_dict)
     23 
     24 torch_script_module = jit.script(model)
     25 torch_script_module.save("./model.pt")