| """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')): |
| |
| 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() |
|
|
