File size: 3,731 Bytes
04a3ade | 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 | """Extract frozen MACE inputs before unused learned preprocessing runs.
The batch plan, MACE call, coordinates and cache layout match multitask_data.
Raising after its existing atomic-input capture avoids downstream preprocessing
that cannot affect the captured raw input. Existing batch-progress files resume.
"""
from pathlib import Path
import sys,gc
sys.path.insert(0,str(Path(__file__).resolve().parent))
from scratch_data import *
common=shared.common
def cache_mace_fast(name,frame):
root=OUT/'cache'/name;root.mkdir(parents=True,exist_ok=True)
if (root/'complete.json').exists():return
from remedi.data_handling.sample import Sample,yield_molecules_collate_fn
from build_author_mace_runtime import build_author_mace_runtime
from train_cosmo import batches
numbers=[torch.tensor(np.asarray(z),dtype=torch.long) for z in frame.atomic_numbers]
positions=[torch.tensor(np.stack(p),dtype=torch.float32) for p in frame.positions]
lengths=np.asarray([len(z) for z in numbers]);ptr=np.r_[0,np.cumsum(lengths)]
np.save(root/'offsets.npy',ptr)
pos=np.lib.format.open_memmap(root/'positions.npy',mode='w+',dtype=np.float32,shape=(int(ptr[-1]),3))
for i,p in enumerate(positions):pos[ptr[i]:ptr[i+1]]=p.numpy()
pos.flush();del pos
mode='r+' if (root/'features.npy').exists() else 'w+'
features=np.lib.format.open_memmap(root/'features.npy',mode=mode,dtype=np.float32,shape=(int(ptr[-1]),4096))
progress=json.loads((root/'progress.json').read_text()) if (root/'progress.json').exists() else {'batches':0}
mace=build_author_mace_runtime().cuda().eval().requires_grad_(False)
captures=[]
class RawCaptured(Exception):
pass
def capture_raw(module, arguments):
captures.append(arguments[0].detach())
raise RawCaptured
hook=mace.preprocessor.atomic_preprocessor.register_forward_pre_hook(capture_raw)
started=time.time()
if name=='bigsoldb':
plan=list(common.batch_indices(np.arange(len(frame)),lengths,max_size=24,atom_cap=2400))
else:plan=list(batches(np.arange(len(frame)),lengths,max_size=64,atom_cap=3600))
status('caching_frozen_mace',corpus=name,total=len(frame),atoms=int(ptr[-1]),bytes=int(ptr[-1])*4096*4,batches=len(plan))
with torch.inference_mode():
for bi,ids in enumerate(plan,1):
if bi<=progress['batches']:continue
batch=yield_molecules_collate_fn([Sample(atomic_positions=positions[i],atomic_numbers=numbers[i],total_charge=float(frame.charge.iloc[i]) if 'charge' in frame else 0.,multiplicity=1.) for i in ids]).to_('cuda')
captures.clear()
try:
mace.preprocessor(batch)
except RawCaptured:
pass
assert len(captures)==1
raw=captures.pop().float().cpu().numpy()
for j,i in enumerate(ids):
v=raw[j,:lengths[i]];assert v.shape==(lengths[i],4096) and np.isfinite(v).all()
features[ptr[i]:ptr[i+1]]=v
if bi%25==0 or bi==len(plan):
features.flush();write(root/'progress.json',dict(batches=bi,total_batches=len(plan)))
status('caching_frozen_mace',corpus=name,batch=bi,batches=len(plan),seconds=time.time()-started)
features.flush();del features
hook.remove();del mace,numbers,positions;gc.collect();torch.cuda.empty_cache()
write(root/'complete.json',dict(rows=len(frame),atoms=int(ptr[-1]),width=4096,dtype='float32',
mace_sha256='fab8b8713c832f31a2a853aaa22fd638be8a369cbf5095e6b3e982a18d10e93a',
raw_capture_only=True,description='Frozen MACE only. Full-precision inputs to learned atomic/geometric preprocessors; no trainable activations cached.'))
|