openfold icon indicating copy to clipboard operation
openfold copied to clipboard

The coordinates of Amino acids are calculated incorrectly after train and using *ckpt in run_pretrained_openfold.py

Open vetmax7 opened this issue 1 year ago • 1 comments
trafficstars

Hello!

Explain me please, how can I use *.ckpt files after train OpenFold for the prediction by run_pretrained_openfold.py? Or perhaps I need to somehow convert this to another format first?

If I use this: --openfold_checkpoint_path /checkpoints/14my.ckpt I got this error message:

raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format( RuntimeError: Error(s) in loading state_dict for AlphaFold: Missing key(s) in state_dict: "aux_heads.tm.linear.weight", "aux_heads.tm.linear.bias".

vetmax7 avatar Mar 05 '24 07:03 vetmax7

I resolved the error by using --config_preset "model_1" instead of model_1_ptm.

However now I have pdb where the coordinates of amino acids are calculated incorrectly. I use "manual optimization mode" because I am testing my own optimizer, the standard Adam generally gives normal results.

Is it a problem with Manual Optimization of Lightning or what?

2024-03-07_18-13-58

Standard Adam and Automatic Optimization:

stdA

vetmax7 avatar Mar 07 '24 08:03 vetmax7

So, I've found a way how to convert *ckpt to my own *npz model. First, in train_openfold.py I added: `def convert_to_pt(ckpt_path, output_path):

checkpoint = torch.load(ckpt_path)
model_state_dict = checkpoint['state_dict']

adjusted_state_dict = {}
for key in model_state_dict.keys():
    adjusted_state_dict[key.replace('model.', "", 1)] = model_state_dict[key]

torch.save(adjusted_state_dict, output_path)
print(f"Converted checkpoint '{ckpt_path}' to PyTorch state dict '{output_path}'.")`

and after if(args.resume_from_ckpt): convert_to_pt(args.resume_from_ckpt, "/npz/model.pt") print("Checkpoint converted to pt format")

Then, I used generated *.pt in convert_of_weights_to_jax.py for creating *.npz

python convert_of_weights_to_jax.py /path_to/model.pt model_1 /path_to_output/model1.npz

I used model_1 instead of model_1_ptm

vetmax7 avatar Apr 08 '24 04:04 vetmax7