File size: 5,017 Bytes
d789132 017b9f5 d789132 017b9f5 d789132 017b9f5 d789132 017b9f5 d789132 017b9f5 d789132 017b9f5 d789132 017b9f5 d789132 017b9f5 d789132 017b9f5 d789132 017b9f5 d789132 | 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 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | import argparse
import json
import time
from pathlib import Path
import numpy as np
import pandas as pd
import torch
import dlmodel as dl
parser = argparse.ArgumentParser()
parser.add_argument('--train', required=True)
parser.add_argument('--output', default='weights')
parser.add_argument('--epochs', type=int, default=12)
parser.add_argument('--holdout', type=float, default=0.1)
parser.add_argument('--split', default='')
parser.add_argument('--alpha', type=float, default=0.7)
parser.add_argument('--drop', type=float, default=0.3)
parser.add_argument('--wd', type=float, default=0.05)
parser.add_argument('--noise', type=float, default=0.1)
parser.add_argument('--crop', type=float, default=0.6)
args = parser.parse_args()
BATCH, LR, CLIP = 64, 2e-3, 1.0
dl.NOISE_P = args.noise
dl.MINKEEP = args.crop
dl.seed_all()
dev = dl.device
cfg = dl.CFG
print('device:', dev)
df = pd.read_parquet(args.train).drop_duplicates(['sequence', 'ec_number'])
vocab_file = Path(__file__).parent / 'ec_vocabulary.json'
if vocab_file.exists():
assert set(df.ec_number) <= set(json.loads(vocab_file.read_text())), 'EC outside the vocabulary'
grouped = df.groupby('sequence', sort=False)['ec_number'].apply(lambda s: sorted(set(s)))
all_seqs, all_ecs = grouped.index.tolist(), grouped.tolist()
ec_list = sorted(df.ec_number.unique())
ec_index = {e: i for i, e in enumerate(ec_list)}
n_classes = len(ec_list)
if args.split:
hold_ids = np.random.RandomState(0).permutation(np.array(json.load(open(args.split))['holdout']))
train_ids = np.setdiff1d(np.arange(len(all_seqs)), hold_ids)
else:
perm = np.random.RandomState(0).permutation(len(all_seqs))
n = int(len(all_seqs) * args.holdout)
hold_ids, train_ids = perm[:n], perm[n:]
n_hold = len(hold_ids)
tr_seqs = [all_seqs[i] for i in train_ids]
tr_labels = [[ec_index[e] for e in all_ecs[i]] for i in train_ids]
tr_lens = [len(s) for s in tr_seqs]
hold_seqs = [all_seqs[i] for i in hold_ids]
hold_ecs = [all_ecs[i] for i in hold_ids]
print(f'{len(tr_seqs)} training sequences, {n_hold} held out, {n_classes} classes')
def make_targets(batch_ids):
rows, cols, vals = [], [], []
for r, i in enumerate(batch_ids):
labs = tr_labels[i]
rows += [r] * len(labs)
cols += labs
vals += [1.0 / len(labs)] * len(labs)
t = torch.zeros(len(batch_ids), n_classes, device=dev)
t[torch.tensor(rows, device=dev), torch.tensor(cols, device=dev)] = torch.tensor(vals, device=dev)
return t
def holdout_l4(P, n=5000):
if not n_hold:
return float('nan')
_, probs = dl.run(P, cfg, hold_seqs[:n], dev)
preds = [ec_list[i] for i in probs.argmax(axis=1)]
return dl.level_scores(preds, hold_ecs[:n])[4]
P = dl.init_params(cfg, n_classes, dev)
S = dl.adam_init(P)
print('parameters:', sum(v.numel() for v in P.values()))
total_steps = len(dl.make_batches(tr_lens, BATCH, True)) * args.epochs
step = 0
best_val, best_epoch, best_P = -1.0, 0, None
for epoch in range(1, args.epochs + 1):
t0, total, count = time.time(), 0.0, 0
for b in dl.make_batches(tr_lens, BATCH, True):
x = dl.make_batch([tr_seqs[i] for i in b], True, dev)
logits, _, C = dl.net_fwd(x, P, cfg, args.drop)
loss, dlogits = dl.cross_entropy_soft_loss_and_grad(logits, make_targets(b))
G = dl.net_bwd(dlogits, P, cfg, C)
norm = dl.adam_step(P, G, S, dl.lr_at(step, total_steps, LR), CLIP, args.wd)
if not np.isfinite(norm):
raise SystemExit('gradient is nan, try a smaller LR in train.py')
step += 1
total += loss.item() * len(b)
count += len(b)
if step % 500 == 0:
print(f' step {step}/{total_steps} loss {total / count:.4f} {time.time() - t0:.0f}s', flush=True)
val = holdout_l4(P)
if n_hold and val > best_val:
best_val, best_epoch = val, epoch
best_P = {k: v.clone() for k, v in P.items()}
print(f'epoch {epoch}/{args.epochs} loss {total / count:.4f} holdout L4 {val:.4f} {time.time() - t0:.0f}s', flush=True)
if best_P is not None:
P = best_P
print(f'using epoch {best_epoch} with holdout L4 {best_val:.4f}')
bank_emb, _ = dl.run(P, cfg, tr_seqs, dev, bs=128, want_probs=False)
bank = dl.make_bank(bank_emb, tr_labels, ec_list)
best_alpha = args.alpha
if n_hold:
Pc = {k: v.cpu() for k, v in P.items()}
torch.set_num_threads(2)
t0 = time.time()
emb, probs = dl.run(Pc, cfg, hold_seqs, 'cpu')
votes = dl.knn_votes(bank, emb)
print(f'CPU inference on {len(hold_seqs)} held out sequences: {time.time() - t0:.0f}s')
best = -1.0
for alpha in (0.0, 0.3, 0.5, 0.7, 1.0):
sc = dl.level_scores(dl.pick(probs, votes, ec_list, alpha), hold_ecs)
print(f'alpha={alpha} ' + ' '.join(f'L{l}={sc[l]:.4f}' for l in (4, 3, 2, 1)))
if sc[4] > best:
best, best_alpha = sc[4], alpha
print('best alpha:', best_alpha)
dl.save(P, cfg, bank_emb, tr_labels, ec_list, args.output, best_alpha)
print('saved to', args.output)
|