File size: 4,041 Bytes
6f1e670 | 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 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 | #!/public/home/scnb9biwet/.conda/envs/model_bio/bin/python
"""
Script to train multievolve neural network models.
Example usage:
conda activate multievolve
p1_train.py \
--experiment-name multievolve_example \
--protein-name example_protein \
--wt-files apex.fasta \
--training-dataset-fname example_dataset.csv \
--wandb-key <key> \
--mode test
"""
import wandb
import argparse
import sys
import matplotlib
matplotlib.use('Agg')
from model.splitters import *
from model.featurizers import *
from model.predictors import *
from model.proposers import *
## LL 2026
import os
os.environ["WANDB_MODE"] = "disabled"
os.environ["WANDB_SILENT"] = "true"
## LL 2026
def parse_args():
"""Parse command line arguments."""
parser = argparse.ArgumentParser(description='Train multievolveneural network models')
parser.add_argument(
'-e',
'--experiment-name',
required=True,
help='Name of the experiment'
)
parser.add_argument(
'-p',
'--protein-name',
required=True,
help='Name of the protein'
)
parser.add_argument(
'-wt',
'--wt-files',
required=True,
help='Comma separated list of paths to the wildtype FASTA files'
)
parser.add_argument(
'-t',
'--training-dataset-fname',
required=True,
help='Path to the training dataset CSV file'
)
parser.add_argument(
'-k',
'--wandb-key',
required=True,
help='WandB API key for authentication'
)
parser.add_argument(
'-m',
'--mode',
required=True,
help='Training method of the experiment, options include: test or standard'
)
args = parser.parse_args()
args.wt_files = [f.strip() for f in args.wt_files.split(',')]
return args
def main():
"""Main function."""
# Parse command line arguments
args = parse_args()
try:
# Login to WandB
#wandb.login(key=args.wandb_key)
# 强制禁用 wandb #LL 2026
import os
os.environ["WANDB_MODE"] = "disabled"
except Exception as e:
print(f"Error logging into WandB: {e}")
sys.exit(1)
# Define variables
experiment_name = args.experiment_name
protein_name = args.protein_name
wt_files = args.wt_files
training_dataset_fname = args.training_dataset_fname
try:
# Define splits
fold_splitter = KFoldProteinSplitter(protein_name, training_dataset_fname, wt_files, csv_has_header=True, use_cache=True, y_scaling=True, val_split=0.15)
splits = fold_splitter.generate_splits(n_splits=5)
except Exception as e:
print(f"Error generating splits: {e}")
sys.exit(1)
try:
# Define features
onehot = OneHotFeaturizer(protein=protein_name, use_cache=True)
features = [onehot]
except Exception as e:
print(f"Error generating features: {e}")
sys.exit(1)
# Define models
models = [Fcn]
if args.mode == 'test':
print("Running in test mode")
sweep_depth = 'test'
search_method = 'test'
elif args.mode == 'standard':
print("Running in standard mode")
sweep_depth = 'standard'
search_method = 'grid'
try:
# Run experiments
print(f"Running experiments for {experiment_name} with {protein_name}...")
run_nn_model_experiments(splits,
features,
models,
experiment_name=experiment_name,
use_cache=True,#LL 2026
sweep_depth=sweep_depth,
search_method=search_method,
show_plots=True, # prevents issue when running script in terminal
)
except Exception as e:
print(f"Error running experiments: {e}")
sys.exit(1)
if __name__ == '__main__':
main()
|