import gc
import math
import os
import shutil
from argparse import Namespace, ArgumentParser, FileType
import torch.nn.functional as F
import wandb
import torch
from sklearn.metrics import roc_auc_score
from torch_geometric.loader import DataListLoader, DataLoader
from tqdm import tqdm
from confidence.dataset import ConfidenceDataset
from utils.training import AverageMeter
torch.multiprocessing.set_sharing_strategy('file_system')
import yaml
from utils.utils import save_yaml_file, get_optimizer_and_scheduler, get_model
parser = ArgumentParser()
parser.add_argument('--config', type=FileType(mode='r'), default=None)
parser.add_argument('--original_model_dir', type=str, default='workdir', help='Path to folder with trained model and hyperparameters')
parser.add_argument('--restart_dir', type=str, default=None, help='')
parser.add_argument('--use_original_model_cache', action='store_true', default=False, help='If this is true, the same dataset as in the original model will be used. Otherwise, the dataset parameters are used.')
parser.add_argument('--data_dir', type=str, default='data/PDBBind_processed/', help='Folder containing original structures')
parser.add_argument('--ckpt', type=str, default='best_model.pt', help='Checkpoint to use inside the folder')
parser.add_argument('--model_save_frequency', type=int, default=0, help='Frequency with which to save the last model. If 0, then only the early stopping criterion best model is saved and overwritten.')
parser.add_argument('--best_model_save_frequency', type=int, default=0, help='Frequency with which to save the best model. If 0, then only the early stopping criterion best model is saved and overwritten.')
parser.add_argument('--run_name', type=str, default='test_confidence', help='')
parser.add_argument('--project', type=str, default='diffdock_confidence', help='')
parser.add_argument('--split_train', type=str, default='data/splits/timesplit_no_lig_overlap_train', help='Path of file defining the split')
parser.add_argument('--split_val', type=str, default='data/splits/timesplit_no_lig_overlap_val', help='Path of file defining the split')
parser.add_argument('--split_test', type=str, default='data/splits/timesplit_test', help='Path of file defining the split')
# Inference parameters for creating the positions and rmsds that the confidence predictor will be trained on.
parser.add_argument('--cache_path', type=str, default='data/cacheNew', help='Folder from where to load/restore cached dataset')
parser.add_argument('--cache_ids_to_combine', nargs='+', type=str, default=None, help='RMSD value below which a prediction is considered a postitive. This can also be multiple cutoffs.')
parser.add_argument('--cache_creation_id', type=int, default=None, help='number of times that inference is run on the full dataset before concatenating it and coming up with the full confidence dataset')
parser.add_argument('--wandb', action='store_true', default=False, help='')
parser.add_argument('--inference_steps', type=int, default=2, help='Number of denoising steps')
parser.add_argument('--samples_per_complex', type=int, default=3, help='')
parser.add_argument('--balance', action='store_true', default=False, help='If this is true than we do not force the samples seen during training to be the same amount of negatives as positives')
parser.add_argument('--rmsd_prediction', action='store_true', default=False, help='')
parser.add_argument('--rmsd_classification_cutoff', nargs='+', type=float, default=2, help='RMSD value below which a prediction is considered a postitive. This can also be multiple cutoffs.')
parser.add_argument('--log_dir', type=str, default='workdir', help='')
parser.add_argument('--main_metric', type=str, default='accuracy', help='Metric to track for early stopping. Mostly [loss, accuracy, ROC AUC]')
parser.add_argument('--main_metric_goal', type=str, default='max', help='Can be [min, max]')
parser.add_argument('--transfer_weights', action='store_true', default=False, help='')
parser.add_argument('--batch_size', type=int, default=5, help='')
parser.add_argument('--lr', type=float, default=1e-3, help='')
parser.add_argument('--w_decay', type=float, default=0.0, help='')
parser.add_argument('--scheduler', type=str, default='plateau', help='')
parser.add_argument('--scheduler_patience', type=int, default=20, help='')
parser.add_argument('--n_epochs', type=int, default=5, help='')
# Dataset
parser.add_argument('--limit_complexes', type=int, default=0, help='')
parser.add_argument('--all_atoms', action='store_true', default=True, help='')
parser.add_argument('--multiplicity', type=int, default=1, help='')
parser.add_argument('--chain_cutoff', type=float, default=10, help='')
parser.add_argument('--receptor_radius', type=float, default=30, help='')
parser.add_argument('--c_alpha_max_neighbors', type=int, default=10, help='')
parser.add_argument('--atom_radius', type=float, default=5, help='')
parser.add_argument('--atom_max_neighbors', type=int, default=8, help='')
parser.add_argument('--matching_popsize', type=int, default=20, help='')
parser.add_argument('--matching_maxiter', type=int, default=20, help='')
parser.add_argument('--max_lig_size', type=int, default=None, help='Maximum number of heavy atoms')
parser.add_argument('--remove_hs', action='store_true', default=False, help='remove Hs')
parser.add_argument('--num_conformers', type=int, default=1, help='')
parser.add_argument('--esm_embeddings_path', type=str, default=None,help='If this is set then the LM embeddings at that path will be used for the receptor features')
parser.add_argument('--no_torsion', action='store_true', default=False, help='')
# Model
parser.add_argument('--num_conv_layers', type=int, default=2, help='Number of interaction layers')
parser.add_argument('--max_radius', type=float, default=5.0, help='Radius cutoff for geometric graph')
parser.add_argument('--scale_by_sigma', action='store_true', default=True, help='Whether to normalise the score')
parser.add_argument('--ns', type=int, default=16, help='Number of hidden features per node of order 0')
parser.add_argument('--nv', type=int, default=4, help='Number of hidden features per node of order >0')
parser.add_argument('--distance_embed_dim', type=int, default=32, help='')
parser.add_argument('--cross_distance_embed_dim', type=int, default=32, help='')
parser.add_argument('--no_batch_norm', action='store_true', default=False, help='If set, it removes the batch norm')
parser.add_argument('--use_second_order_repr', action='store_true', default=False, help='Whether to use only up to first order representations or also second')
parser.add_argument('--cross_max_distance', type=float, default=80, help='')
parser.add_argument('--dynamic_max_cross', action='store_true', default=False, help='')
parser.add_argument('--dropout', type=float, default=0.0, help='MLP dropout')
parser.add_argument('--embedding_type', type=str, default="sinusoidal", help='')
parser.add_argument('--sigma_embed_dim', type=int, default=32, help='')
parser.add_argument('--embedding_scale', type=int, default=10000, help='')
parser.add_argument('--confidence_no_batchnorm', action='store_true', default=False, help='')
parser.add_argument('--confidence_dropout', type=float, default=0.0, help='MLP dropout in confidence readout')
args = parser.parse_args()
if args.config:
config_dict = yaml.load(args.config, Loader=yaml.FullLoader)
arg_dict = args.__dict__
for key, value in config_dict.items():
if isinstance(value, list):
for v in value:
arg_dict[key].append(v)
else:
arg_dict[key] = value
args.config = args.config.name
assert(args.main_metric_goal == 'max' or args.main_metric_goal == 'min')
def train_epoch(model, loader, optimizer, rmsd_prediction):
model.train()
meter = AverageMeter(['confidence_loss'])
for data in tqdm(loader, total=len(loader)):
if device.type == 'cuda' and len(data) % torch.cuda.device_count() == 1 or device.type == 'cpu' and data.num_graphs == 1:
print("Skipping batch of size 1 since otherwise batchnorm would not work.")
optimizer.zero_grad()
try:
pred = model(data)
if rmsd_prediction:
labels = torch.cat([graph.rmsd for graph in data]).to(device) if isinstance(data, list) else data.rmsd
confidence_loss = F.mse_loss(pred, labels)
else:
if isinstance(args.rmsd_classification_cutoff, list):
labels = torch.cat([graph.y_binned for graph in data]).to(device) if isinstance(data, list) else data.y_binned
confidence_loss = F.cross_entropy(pred, labels)
else:
labels = torch.cat([graph.y for graph in data]).to(device) if isinstance(data, list) else data.y
confidence_loss = F.binary_cross_entropy_with_logits(pred, labels)
confidence_loss.backward()
optimizer.step()
meter.add([confidence_loss.cpu().detach()])
except RuntimeError as e:
if 'out of memory' in str(e):
print('| WARNING: ran out of memory, skipping batch')
for p in model.parameters():
if p.grad is not None:
del p.grad # free some memory
torch.cuda.empty_cache()
gc.collect()
continue
else:
raise e
return meter.summary()
def test_epoch(model, loader, rmsd_prediction):
model.eval()
meter = AverageMeter(['loss'], unpooled_metrics=True) if rmsd_prediction else AverageMeter(['confidence_loss', 'accuracy', 'ROC AUC'], unpooled_metrics=True)
all_labels = []
all_affinities = []
for data in tqdm(loader, total=len(loader)):
try:
with torch.no_grad():
pred = model(data)
affinity_loss = torch.tensor(0.0, dtype=torch.float, device=pred[0].device)
accuracy = torch.tensor(0.0, dtype=torch.float, device=pred[0].device)
if rmsd_prediction:
labels = torch.cat([graph.rmsd for graph in data]).to(device) if isinstance(data, list) else data.rmsd
confidence_loss = F.mse_loss(pred, labels)
meter.add([confidence_loss.cpu().detach()])
else:
if isinstance(args.rmsd_classification_cutoff, list):
labels = torch.cat([graph.y_binned for graph in data]).to(device) if isinstance(data,list) else data.y_binned
confidence_loss = F.cross_entropy(pred, labels)
else:
labels = torch.cat([graph.y for graph in data]).to(device) if isinstance(data, list) else data.y
confidence_loss = F.binary_cross_entropy_with_logits(pred, labels)
accuracy = torch.mean((labels == (pred > 0).float()).float())
try:
roc_auc = roc_auc_score(labels.detach().cpu().numpy(), pred.detach().cpu().numpy())
except ValueError as e:
if 'Only one class present in y_true. ROC AUC score is not defined in that case.' in str(e):
roc_auc = 0
else:
raise e
meter.add([confidence_loss.cpu().detach(), accuracy.cpu().detach(), torch.tensor(roc_auc)])
all_labels.append(labels)
except RuntimeError as e:
if 'out of memory' in str(e):
print('| WARNING: ran out of memory, skipping batch')
for p in model.parameters():
if p.grad is not None:
del p.grad # free some memory
torch.cuda.empty_cache()
continue
else:
raise e
all_labels = torch.cat(all_labels)
if rmsd_prediction:
baseline_metric = ((all_labels - all_labels.mean()).abs()).mean()
else:
baseline_metric = all_labels.sum() / len(all_labels)
results = meter.summary()
results.update({'baseline_metric': baseline_metric})
return meter.summary(), baseline_metric
def train(args, model, optimizer, scheduler, train_loader, val_loader, run_dir):
best_val_metric = math.inf if args.main_metric_goal == 'min' else 0
best_epoch = 0
print("Starting training...")
for epoch in range(args.n_epochs):
logs = {}
train_metrics = train_epoch(model, train_loader, optimizer, args.rmsd_prediction)
print("Epoch {}: Training loss {:.4f}".format(epoch, train_metrics['confidence_loss']))
val_metrics, baseline_metric = test_epoch(model, val_loader, args.rmsd_prediction)
if args.rmsd_prediction:
print("Epoch {}: Validation loss {:.4f}".format(epoch, val_metrics['confidence_loss']))
else:
print("Epoch {}: Validation loss {:.4f} accuracy {:.4f}".format(epoch, val_metrics['confidence_loss'], val_metrics['accuracy']))
if args.wandb:
logs.update({'valinf_' + k: v for k, v in val_metrics.items()}, step=epoch + 1)
logs.update({'train_' + k: v for k, v in train_metrics.items()}, step=epoch + 1)
logs.update({'mean_rmsd' if args.rmsd_prediction else 'fraction_positives': baseline_metric,
'current_lr': optimizer.param_groups[0]['lr']})
wandb.log(logs, step=epoch + 1)
if scheduler:
scheduler.step(val_metrics[args.main_metric])
state_dict = model.module.state_dict() if device.type == 'cuda' else model.state_dict()
if args.main_metric_goal == 'min' and val_metrics[args.main_metric] < best_val_metric or \
args.main_metric_goal == 'max' and val_metrics[args.main_metric] > best_val_metric:
best_val_metric = val_metrics[args.main_metric]
best_epoch = epoch
torch.save(state_dict, os.path.join(run_dir, 'best_model.pt'))
if args.model_save_frequency > 0 and (epoch + 1) % args.model_save_frequency == 0:
torch.save(state_dict, os.path.join(run_dir, f'model_epoch{epoch+1}.pt'))
if args.best_model_save_frequency > 0 and (epoch + 1) % args.best_model_save_frequency == 0:
shutil.copyfile(os.path.join(run_dir, 'best_model.pt'), os.path.join(run_dir, f'best_model_epoch{epoch+1}.pt'))
torch.save({
'epoch': epoch,
'model': state_dict,
'optimizer': optimizer.state_dict(),
}, os.path.join(run_dir, 'last_model.pt'))
print("Best Validation accuracy {} on Epoch {}".format(best_val_metric, best_epoch))
def construct_loader_confidence(args, device):
common_args = {'cache_path': args.cache_path, 'original_model_dir': args.original_model_dir, 'device': device,
'inference_steps': args.inference_steps, 'samples_per_complex': args.samples_per_complex,
'limit_complexes': args.limit_complexes, 'all_atoms': args.all_atoms, 'balance': args.balance, 'rmsd_classification_cutoff': args.rmsd_classification_cutoff,
'use_original_model_cache': args.use_original_model_cache, 'cache_creation_id': args.cache_creation_id, "cache_ids_to_combine": args.cache_ids_to_combine}
loader_class = DataListLoader if torch.cuda.is_available() else DataLoader
exception_flag = False
try:
train_dataset = ConfidenceDataset(split="train", args=args, **common_args)
train_loader = loader_class(dataset=train_dataset, batch_size=args.batch_size, shuffle=True)
except Exception as e:
if 'The generated ligand positions with cache_id do not exist:' in str(e):
print("HAPPENING | Encountered the following exception when loading the confidence train dataset:")
print(str(e))
print("HAPPENING | We are still continuing because we want to try to generate the validation dataset if it has not been created yet:")
exception_flag = True
else: raise e
val_dataset = ConfidenceDataset(split="val", args=args, **common_args)
val_loader = loader_class(dataset=val_dataset, batch_size=args.batch_size, shuffle=True)
if exception_flag: raise Exception('We encountered the exception during train dataset loading: ', e)
return train_loader, val_loader
if __name__ == '__main__':
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
with open(f'{args.original_model_dir}/model_parameters.yml') as f:
score_model_args = Namespace(**yaml.full_load(f))
# construct loader
train_loader, val_loader = construct_loader_confidence(args, device)
model = get_model(score_model_args if args.transfer_weights else args, device, t_to_sigma=None, confidence_mode=True)
optimizer, scheduler = get_optimizer_and_scheduler(args, model, scheduler_mode=args.main_metric_goal)
if args.transfer_weights:
print("HAPPENING | Transferring weights from original_model_dir to the new model after using original_model_dir's arguments to construct the new model.")
checkpoint = torch.load(os.path.join(args.original_model_dir,args.ckpt), map_location=device)
model_state_dict = model.state_dict()
transfer_weights_dict = {k: v for k, v in checkpoint.items() if k in list(model_state_dict.keys())}
model_state_dict.update(transfer_weights_dict) # update the layers with the pretrained weights
model.load_state_dict(model_state_dict)
elif args.restart_dir:
dict = torch.load(f'{args.restart_dir}/last_model.pt', map_location=torch.device('cpu'))
model.module.load_state_dict(dict['model'], strict=True)
optimizer.load_state_dict(dict['optimizer'])
print("Restarting from epoch", dict['epoch'])
numel = sum([p.numel() for p in model.parameters()])
print('Model with', numel, 'parameters')
if args.wandb:
wandb.init(
entity='entity',
settings=wandb.Settings(start_method="fork"),
project=args.project,
name=args.run_name,
config=args
)
wandb.log({'numel': numel})
# record parameters
run_dir = os.path.join(args.log_dir, args.run_name)
yaml_file_name = os.path.join(run_dir, 'model_parameters.yml')
save_yaml_file(yaml_file_name, args.__dict__)
args.device = device
train(args, model, optimizer, scheduler, train_loader, val_loader, run_dir)