File size: 6,153 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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | """Random-initialized author architecture, embedding denoising and COSMO heads."""
from pathlib import Path
import sys
sys.path.insert(0,str(Path(__file__).resolve().parent))
from scratch_data import *
from torch import nn
import yaml,math
from threedscriptors.configuration.architecture_config import EncoderDecoderArchitectureConfig
from threedscriptors.data_handling.sample import PreprocessedSample
from threedscriptors.training.noise_scheduler import NoiseModule,ConstantSchedule
from threedscriptors.training.pretraining import atom_denoising_loss
from train_cosmo import loss_terms,batches
CONFIG=dict(epochs=10,probes_per_epoch=5,seed=42026,architecture_source_commit='20df6fe9c61154e351d677cb0f77fd716d00beae',
initialization='Random learned Rem3Di weights; no published or PCQM/COSMO checkpoint loaded',
learning_rate_max=1e-4,weight_decay=1e-3,learning_rate_uncertainty=1e-3,gradient_clip=1.,
noise_sigma=.3,denoising='Exact author atomic embedding denoising loss and gradient paths, on PCQM',
uncertainty='sum(L_t/(2*exp(s_t))+s_t/2); raw denoising loss; initial log variance ln(1088)-0.5 for denoising, -0.5 for COSMO',
uncertainty_initialization='Equivalent initial precision to author-denoising per channel; no modification of raw denoising function',
task_order=['rem3di_denoising','cosmo'],cosmo_head='640 -> 256 SiLU -> 155',
pcqm_batch_size=128,pcqm_atom_cap=5250,cosmo_batch_size=96,cosmo_atom_cap=4000,
epoch_definition='Maximum of complete PCQM and COSMO loader lengths; shorter loader cycles; all retained training molecules seen each epoch',
scheduler='OneCycleLR max_lr=1e-4, total optimizer steps computed from ten exact deterministic epoch plans',
decorrelation=False,combisolv=False,mace='Frozen MACE-POLAR-1-M; raw 4096-channel float32 cache',
normalization='Fresh train-only population invariant mean/std across PCQM and COSMO atoms',
probe='Frozen pooled 640; solute640+solvent640+Temperature_K; exact BigSolDB rows; full75731-row Kumo Large context, 4 estimators, seed42',
probe_schedule='ceil(k * epoch_steps / 5), k=1..5, each epoch; fifty test probes; no test-based selection',
checkpoint_policy='All fifty probe points include all weights, optimizer, scheduler and RNG; verified Hugging Face backup before probing')
def uncertainty(losses,log_variance):
return (losses/(2*log_variance.exp())+log_variance/2).sum()
class ScratchJoint(nn.Module):
def __init__(self,checkpoint=None,device='cuda'):
super().__init__()
checkpoint=Path(checkpoint) if checkpoint is not None else None
architecture_path=checkpoint/'post_training_architecture_config.yaml' if checkpoint is not None else OUT/'architecture.yaml'
architecture=yaml.safe_load(architecture_path.read_text())
architecture['mace_config']=None
cfg=EncoderDecoderArchitectureConfig.model_validate(architecture)
# Constructors initialize every learned tensor. Only data statistics are supplied.
stats_path=checkpoint/'normalization.pt' if checkpoint is not None else DATA/'normalization.pt'
stats=torch.load(stats_path,map_location='cpu',weights_only=True)
self.pre=cfg._build_preprocessor()
self.pre.atomic_preprocessor.double();self.pre.geometric_preprocessor.float()
self.pre.atomic_preprocessor.invariant_normalization.set_stats(stats['mean'],stats['std'],dtype=torch.float64)
self.encoder=cfg._build_encoder().float()
self.denoiser=cfg.decoder_config.build().float()
self.cosmo_head=nn.Sequential(nn.Linear(640,256),nn.SiLU(),nn.Linear(256,155))
self.log_variance=nn.Parameter(torch.tensor([math.log(1088)-.5,-.5],dtype=torch.float32))
self.noise=NoiseModule(ConstantSchedule(.3))
if checkpoint is not None:self.load_parts(Path(checkpoint))
self.to(device)
@property
def parts(self):
return {'encoder.pth':self.encoder,'atomic_preprocessor.pth':self.pre.atomic_preprocessor,
'geometric_preprocessor.pth':self.pre.geometric_preprocessor,'denoising_decoder.pth':self.denoiser,
'cosmo_head.pth':self.cosmo_head}
def load_parts(self,path):
manifest=json.loads((path/'manifest.json').read_text())
for name,module in self.parts.items():
assert sha(path/name)==manifest['files'][name]['sha256'],name
module.load_state_dict(torch.load(path/name,map_location='cpu',weights_only=True),strict=True)
assert sha(path/'task_weighting.pth')==manifest['files']['task_weighting.pth']['sha256']
state=torch.load(path/'task_weighting.pth',map_location='cpu',weights_only=True)
with torch.no_grad():self.log_variance.copy_(state['log_variance'])
def encode(self,sample):
clean=self.pre(sample);descriptor=self.encoder(clean)
z=descriptor if torch.is_tensor(descriptor) else descriptor.flat
return clean,descriptor,z
def denoising(self,sample):
clean,descriptor,z=self.encode(sample)
target=clean.preprocessed_atomic_embeddings.clone()
noisy=self.noise(target,clean.padding_mask)
corrupted=PreprocessedSample(preprocessed_atomic_embeddings=noisy,
initial_pair_representation=clean.initial_pair_representation.detach(),
geometrical_encoding=clean.geometrical_encoding.detach(),
padding_mask=clean.padding_mask,pair_mask=clean.pair_mask)
reconstruction=self.denoiser(corrupted,descriptor)
loss=atom_denoising_loss(target,reconstruction,clean.padding_mask,.3)
return loss,z,target
def cosmo(self,sample,target):
_,_,z=self.encode(sample);pred=self.cosmo_head(z)
loss,parts=loss_terms(pred,target)
return loss,z,pred,parts
def optimizer_for(model):
return torch.optim.AdamW([
dict(params=list(model.pre.parameters())+list(model.encoder.parameters())+list(model.denoiser.parameters())+list(model.cosmo_head.parameters()),lr=CONFIG['learning_rate_max']),
dict(params=[model.log_variance],lr=CONFIG['learning_rate_uncertainty'],weight_decay=0.)],weight_decay=CONFIG['weight_decay'])
|