File size: 2,500 Bytes
32fa850 | 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 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | # /// script
# requires-python = ">=3.13"
# dependencies = ["torch", "timm", "safetensors"]
# ///
"""Verify converted safetensors match original .pt checkpoints."""
import json
import sys
from pathlib import Path
sys.path.insert(0, 'original')
import torch
from safetensors.torch import load_file
from models.anomaly_transformer import get_anomaly_transformer
def verify(dataset: str) -> bool:
"""Verify a single converted checkpoint."""
pt_path = Path(f'{dataset}_parameters.pt')
config_path = Path(dataset) / 'config.json'
safetensors_path = Path(dataset) / 'model.safetensors'
# Load original
original = torch.load(pt_path, map_location='cpu', weights_only=False)
original_sd = original.state_dict()
# Load config and rebuild model
with open(config_path) as f:
config = json.load(f)
model = get_anomaly_transformer(
input_d_data=config['input_d_data'],
output_d_data=config['output_d_data'],
patch_size=config['patch_size'],
d_embed=config['d_embed'],
hidden_dim_rate=config['hidden_dim_rate'],
max_seq_len=config['max_seq_len'],
positional_encoding=config['positional_encoding'],
relative_position_embedding=config['relative_position_embedding'],
transformer_n_layer=config['transformer_n_layer'],
transformer_n_head=config['transformer_n_head'],
dropout=config['dropout'],
)
# Load safetensors weights
saved_sd = load_file(str(safetensors_path))
model.load_state_dict(saved_sd)
loaded_sd = model.state_dict()
# Compare
ok = True
for key in original_sd:
if key not in loaded_sd:
print(f' MISSING: {key}')
ok = False
continue
if not torch.equal(original_sd[key], loaded_sd[key]):
diff = (original_sd[key] - loaded_sd[key]).abs().max().item()
print(f' MISMATCH: {key} (max diff={diff})')
ok = False
extra = set(loaded_sd.keys()) - set(original_sd.keys())
if extra:
print(f' EXTRA keys: {extra}')
ok = False
status = 'OK' if ok else 'FAIL'
print(f'{dataset}: {status}')
return ok
def main() -> None:
"""Verify all converted checkpoints."""
datasets = ['MSL', 'SMAP', 'SWaT', 'WADI']
results = {d: verify(d) for d in datasets}
all_ok = all(results.values())
print(f'\nAll passed: {all_ok}')
if not all_ok:
sys.exit(1)
if __name__ == '__main__':
main()
|