| 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) |
|
|
