Possible to use custom dataloader?

[sophiakrix]

Hi there,

I am working with graphs as a basic data structure and therefore use the Deep Graph Library. In detail, I am using their EdgeDataLoader since it returns the batched data as graphs and messaging blocks, needed for my model (Relational Graph Convolutional Neural Network).

Is it possible to use their dataloader, which is based on the PyTorch dataloader for the train_dataloader ?

My code looks like this:

train_loader = dgl.dataloading.EdgeDataLoader(
        g=heterograph,
        eids=train_eid_dict, 
        block_sampler=sampler,
        batch_size=batch_size,
        negative_sampler=neg_sampler,
        num_workers=4,
        pin_memory=True,
    )

Right now I am getting the following error:

  File ".../LinkPredictHetero.py", line 211, in training_step
    input_nodes, positive_graph, negative_graph, blocks = train_batch
TypeError: cannot unpack non-iterable _EdgeDataLoaderIter object

[awaelchli]

What does your custom dataloader return? If it's a generic object, you can't unfold it like a tuple or list like you showed here.
Do you have an example how you would do it in PyTorch? If it is possible in PyTorch, it's almost guaranteed to work with Lightning too.

[sophiakrix]

Hi @awaelchli , the dataloader is an instance of EdgeCollator from DGL. It is used as follows:

for step, (input_nodes, positive_graph, negative_graph, blocks) in enumerate(train_loader):
    blocks = [b.to(device) if device == th.device('cuda') else b for b in blocks]
    positive_graph = positive_graph.to(device)
    negative_graph = negative_graph.to(device)

    pos_score, neg_score = model(
                    positive_graph=positive_graph,
                    negative_graph=negative_graph,
                    blocks=blocks,
                    h=input_features,
                )
...

And according to the collate function it uses it returns a tuple consisting of the input_nodes, positive_graph, negative_graph, blocks just as I was trying to unpack

[awaelchli]

Oh I see, I think I know. The EdgeDataLoader is not a pytorch dataloader. In fact, they are wrapping a pytorch dataloader. I don't understand why they do this. There is probably a reason but something simple you can try is to just pass in the wrapped dataloader into fit:

trainer.fit(model, train_dataloader=dgl_dataloader.dataloader)

or you return it in your train_dataloader hooks in LightningModule:

def train_dataloader(self):
    dgl_dataloader = ...
    return dgl_dataloader.dataloader

[justusschock]

I just found this :) Adding to @awaelchli : you could also use the wrapper class, but some features of lightning might not work. Those are specifically all the features where we either sanity check or replace values of the loader automatically.

Examples are checks for num_workers or automatic sampler replacement for distributed training. The basic training scheme however should be fine :)