[SOLVED] Dataloader returns only some edge weights

[ajmal]

I love this library! Trying to perform node classification over a graph with weighted edges. I am using stochastic training and therefore blocks. The data loader is dgl.dataloading.DataLoader with a dgl.dataloading.NeighborSampler([4, 4]) sampler. I am trying to pass the edge weights to my model, and I retrieve them as follows during a training loop:

x = blocks[0].srcdata['feature']
weights = blocks[0].edata['weight']
y = blocks[-1].dstdata['label']

With the following forward pass:

    def forward(self, blocks, x, weights):
        h = x
        for l, (layer, block) in enumerate(zip(self.layers, blocks)):
            print(f"Length of weights: {len(weights)}, Number of Edges: {block.number_of_edges()}")
            h = layer(block, h, weights)
            if l != len(self.layers) - 1:
                h = F.relu(h)
                h = self.dropout(h)
        return h

Where each layer is just dglnn.SAGEConv(hid_size, hid_size, 'mean'). A single forward pass for one epoch works with the edge tensor the exact same shape as the number of edges, but a second epoch returns the following error. Am I implementing the sampling incorrectly?

File "/Users/ajmal.aziz/opt/miniconda3/envs/kbricks/lib/python3.10/site-packages/dgl/nn/pytorch/conv/sageconv.py", line 218, in forward
    assert edge_weight.shape[0] == graph.number_of_edges()
AssertionError
Length of weights: 725, Number of Edges: 725
Length of weights: 725, Number of Edges: 185

Any help is really appreciated!

[dyru]

When iterating over blocks, you are using the same edge weights retrieved from the first block, right?

weights = blocks[0].edata['weight']
...
h = layer(block, h, weights)
# change the above line to h = layer(block, h, block.edata['weight'])

[ajmal]

@dyru

block.edata['weight']

That solved my issue - thank you so much! All working now, sorry I was lost in the code.