missing keys in dict when batching custom heterogeneous data

[nabsabraham]

hi @rusty1s! I'm trying to set up a heterogeneous graph of Cell nodes and Gene nodes as below:

from torch_geometric.data import HeteroData
num_cells = xx.shape[0]
num_genes = xx.shape[1]
mydata = HeteroData()
mydata['cell'].x = nn.Embedding(embedding_dim=64, num_embeddings=num_cells)
mydata['gene'].x = nn.Embedding(embedding_dim=64, num_embeddings=num_genes)
mydata['cell'].train_mask = train_mask
mydata['cell'].val_mask = val_mask

mydata['cell'].num_nodes = num_cells
mydata['gene'].num_nodes = num_genes
mydata['cell', 'express', 'gene'].edge_index = edge_index
mydata['cell', 'express', 'gene'].edge_attr = edge_weight
mydata['cell'].y = torch.tensor(obs['age'].values)

I use the HGTLoader but I run into an error in the batching process.

kwargs = {'batch_size': 1024, 'num_workers': 2, 'persistent_workers': True}
mytrain_input_nodes = ('cell', mydata['cell'].train_mask)
myval_input_nodes = ('cell', mydata['cell'].val_mask)

train_loader = HGTLoader(mydata, num_samples=[32] * 2, shuffle=True,
                          input_nodes=mytrain_input_nodes, **kwargs)

out = next(iter(train_loader))

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
[<ipython-input-121-97c50b607f58>](https://localhost:8080/#) in <module>()
----> 1 out = next(iter(train_loader))

2 frames
[/usr/local/lib/python3.7/dist-packages/torch_geometric/loader/utils.py](https://localhost:8080/#) in filter_hetero_data(data, node_dict, row_dict, col_dict, edge_dict, perm_dict)
    146     for node_type in data.node_types:
    147         filter_node_store_(data[node_type], out[node_type],
--> 148                            node_dict[node_type])
    149 
    150     for edge_type in data.edge_types:

KeyError: 'gene'

I tried to debug (below) and found the output of the sampler is missing my gene nodes. Am I missing something in the HeteroData init?

self = train_loader
out = train_loader.sample(indices=[1,2,3])

({'cell': tensor([1, 2, 3])}, {}, {}, {}, 3)

Another related question, when would you use HGTLoader versus a NeighborSampler for heterognns classification?

[rusty1s]

You should make sure that there also exists edges from "gene" to "cell", e.g., via data = transforms.ToUndirected()(data).

[nabsabraham]

awesome, thanks! this fixed it :)