mirror of
https://github.com/aqlaboratory/openfold.git
synced 2026-06-04 12:44:26 +08:00
105 lines
3.5 KiB
Python
105 lines
3.5 KiB
Python
# Copyright 2022 AlQuraishi Laboratory
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
# Converts OpenFold .pt checkpoints into AlphaFold .npz ones, which can then be
|
|
# used to run inference using DeepMind's JAX code.
|
|
import argparse
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from openfold.config import model_config
|
|
from openfold.model.model import AlphaFold
|
|
from openfold.utils.import_weights import (
|
|
Param,
|
|
ParamType,
|
|
generate_translation_dict,
|
|
process_translation_dict,
|
|
import_openfold_weights_
|
|
)
|
|
from openfold.utils.tensor_utils import tree_map
|
|
|
|
|
|
def reshape_fn(of_param, af_weight):
|
|
transformations = {
|
|
ParamType.LinearWeight: lambda w: w.transpose(-1, -2),
|
|
ParamType.LinearWeightMHA: lambda w: w.transpose(-1, -2).reshape(af_weight.shape),
|
|
ParamType.LinearMHAOutputWeight: lambda w: w.transpose(-1, -2).reshape(af_weight.shape),
|
|
ParamType.LinearBiasMHA: lambda w: w.reshape(af_weight.shape),
|
|
ParamType.LinearWeightOPM: lambda w: w.transpose(-1, -2).reshape(af_weight.shape),
|
|
ParamType.Other: lambda w: w,
|
|
}
|
|
|
|
if(of_param.stacked):
|
|
of_weight = torch.stack([torch.Tensor(p) for p in of_param.param])
|
|
else:
|
|
of_weight = torch.Tensor(of_param.param)
|
|
|
|
return transformations[of_param.param_type](of_weight)
|
|
|
|
|
|
def transfer(of_dict, af_weight_template):
|
|
for k in of_dict:
|
|
if(type(of_dict[k]) == dict):
|
|
transfer(of_dict[k], af_weight_template[k])
|
|
else:
|
|
reshaped = reshape_fn(of_dict[k], af_weight_template[k])
|
|
reshaped = reshaped.detach().numpy()
|
|
np.copyto(af_weight_template[k], reshaped)
|
|
|
|
|
|
def main(args):
|
|
d = torch.load(args.of_pt_path)
|
|
|
|
config = model_config(args.config_preset)
|
|
model = AlphaFold(config)
|
|
import_openfold_weights_(model=model, state_dict=d)
|
|
|
|
translation = generate_translation_dict(model, args.config_preset)
|
|
translation = process_translation_dict(translation)
|
|
|
|
af_weight_template = np.load(args.template_npz_path)
|
|
af_weight_template = {k:v for k,v in af_weight_template.items() if k in translation}
|
|
zero = lambda n: n * 0
|
|
af_weight_template = tree_map(zero, af_weight_template, np.ndarray)
|
|
|
|
transfer(translation, af_weight_template)
|
|
|
|
np.savez(args.out_path, **af_weight_template)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"of_pt_path", type=str, help="Path to OpenFold .pt checkpoint file"
|
|
)
|
|
parser.add_argument(
|
|
"config_preset", type=str, help="The corresponding config preset"
|
|
)
|
|
parser.add_argument(
|
|
"out_path", type=str, help="Path for output .npz file"
|
|
)
|
|
parser.add_argument(
|
|
"--template_npz_path",
|
|
type=str,
|
|
default="openfold/resources/params/params_model_1_ptm.npz",
|
|
help="""Path to an AlphaFold checkpoint w/ a superset of the OF
|
|
checkpoint's parameters. params_model_1_ptm.npz always works.
|
|
"""
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
main(args)
|