InMemoryDataset to create training set and test set respectively.

[evantianzg]

Hi,

I am using InMemoryDataset to create my own datasets for graph classification. Generally, a dataset is created through InMemoryDataset and then the training set and test set are split by a train ratio or slicing. My code is as shown below:

class MyDataset(InMemoryDataset):
    def __init__(self, root, transform=None, pre_transform=None, pre_filter=None, force_reload=True):
        super().__init__(root, transform, pre_transform, pre_filter, force_reload=force_reload)
        self.data, self.slices = torch.load(self.processed_paths[0])

    @property
    def raw_file_names(self):
        return ['graph_indicator.csv',
                'graph_labels.csv',
                'graph_A.csv',
                'node_features.csv']

    @property
    def processed_file_names(self):
        return ['data.pt']

    def download(self):
        pass

    def process(self):
        gA_df = pd.read_csv(osp.join(self.raw_dir, 'graph_A.csv'))
        gid_df = pd.read_csv(osp.join(self.raw_dir, 'graph_idx.csv'))
        gl_df = pd.read_csv(osp.join(self.raw_dir, 'graph_labels.csv'), header=None, names=['label'])
        nf_df = pd.read_csv(osp.join(self.raw_dir, 'node_features.csv'))

        edge_index = torch.tensor(gA_df.values.T, dtype=torch.long) 
        x = torch.tensor(nf_df.values, dtype=torch.float)
        y = torch.tensor(gl_df.values.astype(np.int64), dtype=torch.int64)

        i = 0 
        data_list = []
        for subset in gid_df.reset_index().groupby('graph_idx').agg(list).sort_index()['index']:
            subset = torch.tensor(subset, dtype=torch.long) 
            sub_edge_index = subgraph(subset, edge_index, relabel_nodes=True)[0]
            data = Data(edge_index=sub_edge_index, 
                        x=x[subset], 
                        y=y[i])
            i += 1
            data_list.append(data)

        data, slices = self.collate(data_list)
        torch.save((data, slices), self.processed_paths[0])

def split_dataset_by_ratio(dataset, train_ratio, test_datasets):
    labels = [data.y.item() for data in dataset]
    train_dataset, test_dataset = train_test_split(dataset, train_size=train_ratio, stratify=labels, random_state=2024)
    test_datasets.append(test_dataset)
    return train_dataset, test_datasets

Now I need to specify the nodes for training set and test set. So I manually prepared the corresponding csv files for training set and test set. That is: tr-graph_indicator.csv, tr-graph_labels.csv, tr-graph_A.csv, and tr-node_features.csv files for training set, and te-graph_indicator.csv, te-graph_labels.csv, te-graph_A.csv, and te-node_features.csv files for test set.

Now I need to create the training set and test set using InMemoryDataset, respectively. I suppose the simplest way is to define two InMemoryDataset classes (one for training set and another for test set). But I believe it is a stupid method and I think there should be more elegant ways.

Any suggestions are appreciated🙏

[rusty1s]

Why can't you simply do post-filtering and only keep one single dataset as you do in your example? I don't think you need to have two InMemoryDataset variants. Alternatively, you can specify a split attribute as part of your dataset class, and then load in the correct data. This is what we commonly do as well, e.g., here: https://github.com/pyg-team/pytorch_geometric/blob/master/torch_geometric/datasets/gdelt.py

[evantianzg]

Hi, sorry for the late reply. I am just back from a long break.

Yes I have tried the post-filtering and it works fine. I think the split attribute is more elegant.