File size: 2,682 Bytes
d346d4e | 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 | """Train only on the provided training split. No validation labels are used."""
import argparse
import json
from pathlib import Path
import numpy as np
import pandas as pd
from model import FEATURE_NAMES, features
# def train(train_path, output):
# frame = pd.read_parquet(train_path)
# x = features(frame[['smiles', 'sequence']].to_dict('records'))
# y = frame.log10_value.to_numpy(float)
# mean, scale = x.mean(axis=0), x.std(axis=0)
# scale[scale < 1e-10] = 1
# x = (x - mean) / scale
# intercept = float(y.mean())
# coef = np.linalg.solve(x.T @ x + 10 * np.eye(x.shape[1]), x.T @ (y - intercept))
# out = Path(output)
# out.parent.mkdir(parents=True, exist_ok=True)
# out.write_text(json.dumps({'feature_names': FEATURE_NAMES, 'mean': mean.tolist(),
# 'scale': scale.tolist(), 'coef': coef.tolist(), 'intercept': intercept}, indent=2))
# print(f'Trained baseline on {len(frame)} training rows; saved {out}')
def train(train_path, output):
import torch # inline so this stays a drop-in replacement
frame = pd.read_parquet(train_path)
x = features(frame[['smiles', 'sequence']].to_dict('records'))
y = frame.log10_value.to_numpy(float)
mean, scale = x.mean(axis=0), x.std(axis=0)
scale[scale < 1e-10] = 1
x = (x - mean) / scale
intercept = float(y.mean())
torch.manual_seed(0)
xt = torch.tensor(x, dtype=torch.float32)
yt = torch.tensor(y - intercept, dtype=torch.float32)[:, None]
w = torch.zeros(x.shape[1], 1, requires_grad=True)
opt = torch.optim.AdamW([w], lr=1e-2, weight_decay=1e-3)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=500)
for _ in range(500):
perm = torch.randperm(len(xt))
for i in range(0, len(xt), 1024):
b = perm[i:i + 1024]
opt.zero_grad()
torch.nn.functional.huber_loss(xt[b] @ w, yt[b]).backward()
opt.step()
sched.step()
coef = w.detach().numpy().ravel().astype(float)
out = Path(output)
out.parent.mkdir(parents=True, exist_ok=True)
out.write_text(json.dumps({'feature_names': FEATURE_NAMES, 'mean': mean.tolist(),
'scale': scale.tolist(), 'coef': coef.tolist(), 'intercept': intercept}, indent=2))
print(f'Trained baseline on {len(frame)} training rows; saved {out}')
if __name__ == '__main__':
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--train', default='../data/public/train.parquet')
parser.add_argument('--output', default='weights/baseline.json')
args = parser.parse_args()
train(args.train, args.output)
|