File size: 2,028 Bytes
fd684da | 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 | """Export inference states without optimizer, paths, or experiment tracker metadata."""
import argparse
import hashlib
import json
from pathlib import Path
import sys
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import torch
from omegaconf import OmegaConf
from models import Encoder, Decoder, ARModel
def sha256(path):
digest=hashlib.sha256()
with open(path,'rb') as f:
for block in iter(lambda:f.read(8*1024*1024),b''): digest.update(block)
return digest.hexdigest()
def export(source, config, output):
source_state=torch.load(source, map_location='cpu', weights_only=True, mmap=True)
c=OmegaConf.load(config)
states={key:{name.removeprefix('_orig_mod.'): tensor for name,tensor in source_state[key].items()}
for key in ('enc','dec','ar_model','ar_model_ema')}
for cls,key in ((Encoder,'enc'),(Decoder,'dec'),(ARModel,'ar_model'),(ARModel,'ar_model_ema')):
# Assignment checks all keys/shapes without allocating a second model's storage.
with torch.device('meta'): model=cls(c)
model.load_state_dict(states[key],strict=True,assign=True)
del model
out=Path(output); out.parent.mkdir(parents=True,exist_ok=True)
if out.exists(): raise FileExistsError(out)
torch.save(dict(format_version=1,**states),out)
saved=torch.load(out,map_location='cpu',weights_only=True,mmap=True)
for key,state in states.items():
for name,tensor in state.items():
if not torch.equal(tensor,saved[key][name]): raise AssertionError((key,name))
return dict(filename=out.name,sha256=sha256(out),bytes=out.stat().st_size)
def main():
p=argparse.ArgumentParser(description=__doc__)
p.add_argument('--source',required=True); p.add_argument('--config',required=True); p.add_argument('--output',required=True)
a=p.parse_args(); result=export(a.source,a.config,a.output)
Path(a.output+'.json').write_text(json.dumps(result,indent=2)+'\n'); print(json.dumps(result))
if __name__=='__main__': main()
|