import os from pathlib import Path import torch import argparse parser = argparse.ArgumentParser() args = parser.parse_args() def prune_it(p, keep_only_ema=True): print(f"prunin' in path: {p}") size_initial = os.path.getsize(p) nsd = dict() sd = torch.load(p, map_location="cpu") print(sd.keys()) for k in sd.keys(): if k != "optimizer_states": nsd[k] = sd[k] else: print(f"removing optimizer states for path {p}") if "global_step" in sd: print(f"This is global step {sd['global_step']}.") if keep_only_ema: sd = nsd["state_dict"].copy() # infer ema keys ema_keys = {k: "model_ema." + k[6:].replace(".", "") for k in sd.keys() if k.startswith('model.')} new_sd = dict() for k in sd: if k in ema_keys: print(k, ema_keys[k]) new_sd[k] = sd[ema_keys[k]] elif not k.startswith("model_ema.") or k in ["model_ema.num_updates", "model_ema.decay"]: new_sd[k] = sd[k] assert len(new_sd) == len(sd) - len(ema_keys) nsd["state_dict"] = new_sd else: sd = nsd['state_dict'].copy() new_sd = dict() for k in sd: new_sd[k] = sd[k] nsd['state_dict'] = new_sd fn = f"{os.path.splitext(p)[0]}-pruned.ckpt" if not keep_only_ema else f"{os.path.splitext(p)[0]}-ema-pruned.ckpt" print(f"saving pruned checkpoint at: {fn}") torch.save(nsd, fn) newsize = os.path.getsize(fn) MSG = f"New ckpt size: {newsize*1e-9:.2f} GB. " + \ f"Saved {(size_initial - newsize)*1e-9:.2f} GB by removing optimizer states" if keep_only_ema: MSG += " and non-EMA weights" print(MSG) if __name__ == "__main__": prune_it('wd-v1-2-full-ema.ckpt')