How to properly construct dgl.block for training？

[BearBiscuit05]

I encountered a problem when building my own dataset. My data is separated from the graph structure and graph features, so I want to store only the graph structure information in the block, and then add the corresponding features when putting it into the model to complete the training. But it seems that I can’t get the correct result, and there are some errors, I want to know what’s wrong with my code.

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchmetrics.functional as MF
import dgl
import dgl.nn as dglnn
from dgl.heterograph import DGLBlock
from torch.utils.data import Dataset, DataLoader
class SAGE(nn.Module):
    def __init__(self, in_size, hid_size, out_size):
        super().__init__()
        self.layers = nn.ModuleList()
        self.layers.append(dglnn.SAGEConv(in_size, hid_size, 'mean'))
        self.layers.append(dglnn.SAGEConv(hid_size, out_size, 'mean'))
        self.dropout = nn.Dropout(0.5)

    def forward(self, blocks, x):
        h = x
        for l, (layer, block) in enumerate(zip(self.layers, blocks)):
            h = layer(block, h)
            if l != len(self.layers) - 1:
                h = F.relu(h)
                h = self.dropout(h)
        return h

def create_dgl_block(data, num_src_nodes, num_dst_nodes):
    row, col = data
    gidx = dgl.heterograph_index.create_unitgraph_from_coo(2, num_src_nodes, num_dst_nodes, row, col, 'coo')
    g = DGLBlock(gidx, (['_N'], ['_N']), ['_E'])
    return g

if __name__ == '__main__':
    model = SAGE(100, 128, 4)
    print('Training...')
    blocks = []
    row = torch.tensor([0, 0, 0])
    col = torch.tensor([1, 2, 3])
    block = create_dgl_block((row,col),3,1)
    blocks.append(block)  

    row = torch.tensor([1, 1, 2, 2, 3, 3])
    col = torch.tensor([4, 5, 6, 7, 8, 9])
    block = create_dgl_block((row,col),6,3)
    blocks.insert(0,block)  
    random_feat = torch.rand((6, 100))
    opt = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=5e-4)
    
    model.train()
    y = torch.tensor([1])
    y_hat = model(blocks, random_feat)
    loss = F.cross_entropy(y_hat, y)
    opt.zero_grad()
    loss.backward()
    opt.step()

[mufeili]

Even in that case, it might be a better idea to generate DGL blocks with built-in samplers and dataloaders. You can then manually insert node features and labels.