extract_checkpoint_mean_std.py

import torch
import safetensors.torch
import sys

def dump_keys(parent, suffix=''):
    for k in sorted(parent.keys()):
        if isinstance(parent[k], torch.Tensor):
            print(f'{suffix}{k} {list(parent[k].shape)} mean={torch.mean(parent[k]):.3g} std={torch.std(parent[k]):.3g}')
        else:
            dump_keys(parent[k], f'{suffix}{k}.')

if __name__ == '__main__':
    fn = sys.argv[1]
    if fn.endswith('.safetensors'):
        checkpoint = safetensors.torch.load_file(fn)
    else:
        checkpoint = torch.load(fn)
    dump_keys(checkpoint)