How to enable advanced mini-batching behaviour

[inesochoa]

Hi there,

I have written my own custom data and dataset objects. I am using PyG dataloader but the output I am getting is not batched. I am pasting the relevant classes below, but in short, I am getting a CustomDataset with a single graph, i.e. there is a single x tensor with dimensions [N,t,f] rather than multiple tensors with dimension [batch_size * t, f] (where N is the total number of "graphs" in the sample, t is the number of nodes per graph and f is the number of node features).

How I am creating the dataset and dataloader:

from torch_geometric.data import Data
from torch_geometric.loader import DataLoader
all_ds = CustomDataset(name,tfeatures)
dl = DataLoader(all_ds, batch_size = batchSize, drop_last=True, num_workers=0, shuffle=False, pin_memory=False)

And the corresponding custom data objects (I have a few more attributes in MyData that I left out for simplicity):

from torch_geometric.data import Data, Dataset
class CustomDataset(Dataset):
    def __init__(self,name,tfeatures):
        # load data
        f = h5py.File(name,"r")
        self.y = torch.tensor(f["b"][:]["label"]).view(-1,1).long()
        self.x = torch.tensor(rfn.structured_to_unstructured(f["a"][:][tfeatures]),dtype=torch.float32)
        self.edge_index = single_fully_connected(0,t,transpose=True)
        self.data = MyData(x=self.x, y=self.y, edge_index=self.edge_index)
        f.close()
    
    def __len__(self):
        return self.x.shape[0]
    
class MyData(Data): 
    def __init__(self, x=None, y=None, edge_index=None):
        super().__init__()
        if x is not None: self.x = x
        if y is not None: self.y = y
        if edge_index is not None: self.edge_index = edge_index
    def __len__(self):
        return len(self.x.shape[0])

I am guessing I need to override some of the class functions (maybe collate_fn) but it's not clear to me how.

Many thanks in advance!

[wsad1]

The x tensor in CustomDataset is of shape [N,t,f] , so DataLoader can't reshape that into [batch_size * t, f]. Esentially the DataLoader thinks you have a single graph and its simply returning that.

What you could do is pass it a list of CustomDataset of size N with each containing a x tensor of shape [t,f], then it'll return a dataset where x is of shape [bathc_size*t, f].
See this example for reference.

[inesochoa]

Thanks for the reply! I have implemented your suggestion as:

func = map(lambda a,b,c: MyData(a,b,c), x, y, edge_index)
ds = list(func)

such that I get a list of N elements in the format: MyData(x=[t, f], y=[1], edge_index=[2, t*(t-1)]) and indeed this seems to work correctly when it gets passed to the dataloader.

However, the problem with this approach is that it's much slower; is there any better way to do this (loading the dataset and batching it correctly)?

Thanks,

[wsad1]

Usually training (forward + backward) pass is the bottleneck and not dataloading. If you think data loading is the bottleneck can you try setting num_workers in DataLoader to maybe 4 or 8, you might have to experiment to set the right value, this should speed up data loading.