How to pass the protein pdb file to the released model weight: (mc_gearnet_edge.pth)

[Tizzzzy]

Hi author,
Right now I download the pretrain weight from Multiview Contrast. I want to load this weight, and then pass a pdb file to the model. In this case, I can get the latent representation of the protein pdb file from the model.

Now I think I can successfully load the released pretrained weight, however, I am stuck at passing the pdb file to the model. Here is how I pass it:

def extract_representation(model, protein_structure_path):
 protein = Protein.from_pdb(protein_structure_path)
 
 _protein = Protein.pack([protein])

 input_feature = protein.atom2graph
 
 with torch.no_grad():
 representation = model(_protein, input_feature)
 
 # Return the output, which is the protein's representation
 return representation

I think GearNet model takes two parameter: 1. a Graph object, 2. input (not sure what this is). However, it seems like either _protein is wrong or input_feature is wrong. With this code, I am getting this error:

Traceback (most recent call last):
 File "/content/GearNet-main/script/test.py", line 62, in <module>
 main()
 File "/content/GearNet-main/script/test.py", line 58, in main
 representation = extract_representation(model, args.pdb)
 File "/content/GearNet-main/script/test.py", line 43, in extract_representation
 representation = model(_protein, input_feature)
 File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
 return self._call_impl(*args, **kwargs)
 File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1541, in _call_impl
 return forward_call(*args, **kwargs)
 File "/content/GearNet-main/gearnet/model.py", line 115, in forward
 edge_hidden = self.edge_layers[i](line_graph, edge_hidden)
 File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
 return self._call_impl(*args, **kwargs)
 File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1541, in _call_impl
 return forward_call(*args, **kwargs)
 File "/content/GearNet-main/gearnet/layer.py", line 132, in forward
 update = self.aggregate(graph, message)
 File "/content/GearNet-main/gearnet/layer.py", line 118, in aggregate
 update = update.view(graph.num_node, self.num_relation * self.input_dim)
RuntimeError: shape '[4944, 472]' is invalid for input of size 711936

Can you please help me on how to pass the pdb file to the model.
Thank you

[mpedraza98]

I think your mistake is that you are using atom2graph as input feature, when the packed protein object has the attribute residue_feature or node_feature Depending if you are working with residue or atom view.
Doing something like the code below worked for me

protein_model = graph_construction_model(packed_protein)
representation = gearnet_edge(graph=protein_model, input=protein_model.residue_feature.float())