Module train

Module: train

Finetunes a CheckAMG-PST model to predict whether proteins are (A) viral and (B) auxiliary, for use with de-novo. Use this module to update CheckAMG-PST models or to finetune the model with your own data labeled for viral/non-viral and AVG/non-AVG status.

A model trained here is supplied to de-novo via --model-ckpt together with one of --train-data-file, --train-embed-file, or (--train-index-file and --train-labels-file), instead of the default database.

Basic usage

checkamg train \
  -i train.graphfmt.h5 \
  -mc pretrained.ckpt \
  -o train_output

Parameters

Input train/test data

Argument	Description
-i, --train-file	Graph-formatted .h5 training data (required).
-te, --test-file	Graph-formatted .h5 test data. If omitted, no test set is used.
-mc, --model-ckpt	Pretrained PST checkpoint (.ckpt) to finetune (required).

Outputs

Argument	Default	Description
-o, --output	required	Directory for logs, checkpoints, and the trained model.
-s, --save-train-embed / --no-...	off	Store training protein embeddings in --output. If off, they are not computed after training.

Training parameters

Argument	Default	Description
-b, --batch-size	16	Batch size in number of scaffolds.
--lr	1e-3	Learning rate, range (0.0, 0.1].
-e, --max-epochs	25	Maximum number of epochs.
--dropout	0.0	Dropout, range [0.0, 1.0].
--margin	1.0	Triplet loss margin.
--knn	5	Nearest neighbors for majority-voting class-label prediction.
--max-ap-pairs	250000	Maximum anchor-positive pairs considered at a time.
--positive-mining-strategy	esm_easy	esm_easy (nearest same-class protein in input ESM space), hard (farthest same-class protein in PST space), or all (every same-class pair).
--negative-mining-strategy	semihard	hard, semihard, or both.
--opposite-class-negatives / --no-...	on	Restrict negative sampling, within the hard/semihard candidate set, to the most-opposite viral/AVG class per anchor (priority: flip viral, then flip AVG, then flip both).
--class-weighting	class_freq	class_freq (inverse weighted frequency of class labels) or pair_freq (inverse weighted frequency of same-class pairs).
--context-size	4	Number of proteins up/downstream of an AVG used to align short-context views with the long-context embedding.

Resources, other

Argument	Default	Description
-a, --accelerator	auto	cpu, gpu, or auto.
-x, --devices	1	Number of devices. On CPU this is cores/threads; multi-GPU is not supported.
-m, --mem	80% of available	Memory limit in GB.
--verbose / --no-verbose	off	Per-step logging in the progress bar instead of per-epoch only.
--seed	111	Random seed.

Reproducibility

The notebooks used to build the training/test datasets and to train the released CheckAMG-PST model are described on the Reproducibility page.