File size: 7,988 Bytes
c2767f4 | 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 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 | from argparse import Namespace
from functools import partial
from pathlib import Path
import torch
import yaml
from torch_geometric.nn.data_parallel import DataParallel
from onescience.utils.diffdock.diffusion_utils import get_timestep_embedding, t_to_sigma as t_to_sigma_compl
from onescience.utils.diffdock.utils import ExponentialMovingAverage
from .aa_model import AAModel
from .cg_model import CGModel
from .old_aa_model import AAOldModel
_LM_EMBEDDING_KEYS = (
"moad_esm_embeddings_path",
"pdbbind_esm_embeddings_path",
"pdbsidechain_esm_embeddings_path",
"esm_embeddings_path",
"esm_embeddings_model",
)
def load_model_args(model_dir):
model_dir = Path(model_dir)
config_path = model_dir / "model_parameters.yml"
with config_path.open("r", encoding="utf-8") as handle:
config = yaml.full_load(handle) or {}
return Namespace(**config)
def model_uses_lm_embeddings(model_args):
return any(getattr(model_args, key, None) is not None for key in _LM_EMBEDDING_KEYS)
def _has_arg(args, name):
try:
return name in args
except TypeError:
return hasattr(args, name)
def _get_arg(args, name, default=None):
if isinstance(args, dict):
return args.get(name, default)
if _has_arg(args, name):
return getattr(args, name)
return default
def get_model(args, device, t_to_sigma, no_parallel=False, confidence_mode=False, old=False):
timestep_emb_func = get_timestep_embedding(
embedding_type=_get_arg(args, "embedding_type", "sinusoidal"),
embedding_dim=args.sigma_embed_dim,
embedding_scale=_get_arg(args, "embedding_scale", 10000),
)
all_atoms = _get_arg(args, "all_atoms", False)
if old and not all_atoms:
raise NotImplementedError(
"The old coarse-grained DiffDock model path is not migrated yet. "
"Use old=True only with all_atoms=True, or migrate old_cg_model.py first."
)
lm_embedding_type = None
if (
_get_arg(args, "moad_esm_embeddings_path") is not None
or _get_arg(args, "pdbbind_esm_embeddings_path") is not None
or _get_arg(args, "pdbsidechain_esm_embeddings_path") is not None
or _get_arg(args, "esm_embeddings_path") is not None
):
lm_embedding_type = "precomputed"
if _get_arg(args, "esm_embeddings_model") is not None:
lm_embedding_type = args.esm_embeddings_model
if old:
model_class = AAOldModel
elif all_atoms:
model_class = AAModel
else:
model_class = CGModel
model_kwargs = dict(
t_to_sigma=t_to_sigma,
device=device,
no_torsion=args.no_torsion,
timestep_emb_func=timestep_emb_func,
num_conv_layers=args.num_conv_layers,
lig_max_radius=args.max_radius,
scale_by_sigma=args.scale_by_sigma,
sigma_embed_dim=args.sigma_embed_dim,
norm_by_sigma=_get_arg(args, "norm_by_sigma", False),
ns=args.ns,
nv=args.nv,
distance_embed_dim=args.distance_embed_dim,
cross_distance_embed_dim=args.cross_distance_embed_dim,
batch_norm=not args.no_batch_norm,
dropout=args.dropout,
use_second_order_repr=args.use_second_order_repr,
cross_max_distance=args.cross_max_distance,
dynamic_max_cross=args.dynamic_max_cross,
smooth_edges=_get_arg(args, "smooth_edges", False),
odd_parity=_get_arg(args, "odd_parity", False),
lm_embedding_type=lm_embedding_type,
confidence_mode=confidence_mode,
confidence_dropout=_get_arg(args, "confidence_dropout", 0.0),
confidence_no_batchnorm=_get_arg(args, "confidence_no_batchnorm", False),
affinity_prediction=_get_arg(args, "affinity_prediction", False),
parallel=_get_arg(args, "parallel", 1),
num_confidence_outputs=(
len(args.rmsd_classification_cutoff) + 1
if isinstance(_get_arg(args, "rmsd_classification_cutoff"), list)
else 1
),
atom_num_confidence_outputs=(
len(args.atom_rmsd_classification_cutoff) + 1
if isinstance(_get_arg(args, "atom_rmsd_classification_cutoff"), list)
else 1
),
parallel_aggregators=_get_arg(args, "parallel_aggregators", ""),
fixed_center_conv=not _get_arg(args, "not_fixed_center_conv", False),
no_aminoacid_identities=_get_arg(args, "no_aminoacid_identities", False),
include_miscellaneous_atoms=_get_arg(args, "include_miscellaneous_atoms", False),
sh_lmax=_get_arg(args, "sh_lmax", 2),
differentiate_convolutions=not _get_arg(args, "no_differentiate_convolutions", False),
tp_weights_layers=_get_arg(args, "tp_weights_layers", 2),
num_prot_emb_layers=_get_arg(args, "num_prot_emb_layers", 0),
reduce_pseudoscalars=_get_arg(args, "reduce_pseudoscalars", False),
embed_also_ligand=_get_arg(args, "embed_also_ligand", False),
atom_confidence=_get_arg(args, "atom_confidence_loss_weight", 0.0) > 0.0,
sidechain_pred=(
(_has_arg(args, "sidechain_loss_weight") and args.sidechain_loss_weight > 0)
or (_has_arg(args, "backbone_loss_weight") and args.backbone_loss_weight > 0)
),
depthwise_convolution=_get_arg(args, "depthwise_convolution", False),
)
if model_class is AAModel:
model_kwargs["crop_beyond"] = _get_arg(args, "crop_beyond", None)
elif model_class is AAOldModel:
for key in (
"atom_num_confidence_outputs",
"differentiate_convolutions",
"tp_weights_layers",
"num_prot_emb_layers",
"reduce_pseudoscalars",
"embed_also_ligand",
"atom_confidence",
"sidechain_pred",
"depthwise_convolution",
):
model_kwargs.pop(key, None)
model_kwargs["lm_embedding_type"] = (
"esm" if _get_arg(args, "esm_embeddings_path") is not None else None
)
model_kwargs["use_old_atom_encoder"] = _get_arg(args, "use_old_atom_encoder", True)
model = model_class(**model_kwargs)
if device.type == "cuda" and not no_parallel and _get_arg(args, "dataset") != "torsional":
model = DataParallel(model)
model.to(device)
return model
def build_score_model(
model_args,
device,
no_parallel=False,
confidence_mode=False,
old=False,
):
t_to_sigma = partial(t_to_sigma_compl, args=model_args)
model = get_model(
model_args,
device,
t_to_sigma=t_to_sigma,
no_parallel=no_parallel,
confidence_mode=confidence_mode,
old=old,
)
return model, t_to_sigma
def load_score_model(
model_dir,
ckpt,
device=None,
no_parallel=True,
confidence_mode=False,
old=False,
strict=True,
):
if device is None:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_args = load_model_args(model_dir)
model, t_to_sigma = build_score_model(
model_args=model_args,
device=device,
no_parallel=no_parallel,
confidence_mode=confidence_mode,
old=old,
)
checkpoint_path = Path(model_dir) / ckpt
state_dict = torch.load(checkpoint_path, map_location=torch.device("cpu"))
if isinstance(state_dict, dict) and "model" in state_dict and "optimizer" in state_dict:
model.load_state_dict(state_dict["model"], strict=strict)
if "ema_weights" in state_dict and getattr(model_args, "ema_rate", None) is not None:
ema_weights = ExponentialMovingAverage(model.parameters(), decay=model_args.ema_rate)
ema_weights.load_state_dict(state_dict["ema_weights"], device=device)
ema_weights.copy_to(model.parameters())
else:
model.load_state_dict(state_dict, strict=strict)
model = model.to(device)
model.eval()
return model, model_args, t_to_sigma
|