OpenFold / scripts /convert_of_weights_to_jax.py

Upload folder using huggingface_hub

3c52512 verified 3 months ago

3.63 kB

# 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 onescience.configs.bio.openfold.config import model_config
from onescience.models.openfold.model import AlphaFold
from onescience.utils.openfold.import_weights import (
    Param, 
    ParamType, 
    generate_translation_dict,
    process_translation_dict,
    import_openfold_weights_,
)
from onescience.utils.openfold.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)

	# 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 onescience.configs.bio.openfold.config import model_config
	from onescience.models.openfold.model import AlphaFold
	from onescience.utils.openfold.import_weights import (
	Param,
	ParamType,
	generate_translation_dict,
	process_translation_dict,
	import_openfold_weights_,
	)
	from onescience.utils.openfold.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)