Missing data.x after converting graph to pytorch_geometric Data object

[kamurani]

Hey everyone,

I have been using a modified ProteinGraphDataset via a pyg DataLoader to supply graphs to a GCNConv layer in a neural network model. However, I am not sure how to retrieve the node_features element of the pytorch geometric Data class. According to their documentation, this is stored by default as data.x. When I print out what my loader is supplying, I get this object:

DataBatch(edge_index=[2, 45], node_id=[4], coords=[4], name=[4], dist_mat=[4], num_nodes=61, batch=[61], ptr=[5])

I've looked through the convert_nx_to_pyg_data func in graphein.ml.conversion and notice there's a loop "adding node features" but looks like it's indexing these using the individual feature's name as opposed to keeping them all in one tensor. That's just from what i've understood, maybe I'm missing something.

Everything else I need for the GCNConv looks good.

Any ideas how to get this to work with a model that's expecting to do this:

pro_x, pro_edge_index, pro_batch = pro_data.x, pro_data.edge_index, pro_data.batch
x = self.pro_conv(pro_x, pro_edge_index) 
# The issue is here ^ , where pro_x is `None` 
x = ....


... 

return out

Thanks in advance for anyone's help! Really appreciate it

[a-r-j]

We keep attributes separate by design (e.g. you may wish to do different things to them, such as treating coordintates in a more principled way in an EGNN).

So, what you can do to solve your problem is simply concatenate the features you want

def forward(batch: Batch):
    x = torch.cat([batch.amino_acid_one_hot, batch.some_other_feature], dim=1)
    x = self.conv(x, batch.edge_index)

If you want to do this faster (and are using pytorch lightning modules), you can do the pre-processing in on_after_batch_transfer