Creating Large Datasets

[anthonysirico]

Prior to this, I used InMemoryDataset, now I would like to use the Large Dataset method. Two things:

1. Did I set this up correctly?
2. The .get() method is not working, I get the error TypeError: TMSDataset.get() missing 1 required positional argument: 'self'. The goal when I run get() is to load all 32,000+ graphs.

import os.path as osp
import pickle
from abc import ABC
from tqdm import tqdm
import pandas as pd
import torch
from torch_geometric.data import Dataset
from torch_geometric.utils import from_networkx
import networkx as nx
from sklearn.preprocessing import LabelEncoder


def _get_node_features(networkx_graph, node_labels):
    nl = node_labels[0:len(networkx_graph.nodes)]
    nl = dict(enumerate(nl))
    nx.set_node_attributes(networkx_graph, nl, 'feature')
    return networkx_graph


class TMSDataset(Dataset, ABC):
    def __init__(self, root, transform=None, pre_transform=None, pre_filter=None):
        super(TMSDataset, self).__init__(root, transform, pre_transform, pre_filter)
        self.data_path = 'path'

    @property
    def raw_file_names(self):
        return ['networkx_graphs.pickle', 'nl.json', 'graph_features.csv']

    @property
    def processed_file_names(self):
        return ['tms_graph_{}.pt'.format(i) for i in range(32612)]

    def download(self):
        pass

    def process(self):
        idx = 0
        networkx_graphs = pickle.load(open(self.raw_paths[0], 'rb'))
        node_labels = pd.read_json(self.raw_paths[1])
        graph_features = pd.read_csv(self.raw_paths[2])
        encoder = LabelEncoder()
        encoded_node_labels = node_labels.apply(encoder.fit_transform, axis=1)

        for index in tqdm(range(len(networkx_graphs))):
            nx_graph = networkx_graphs[index]
            nx_graph = _get_node_features(nx_graph, encoded_node_labels[index])
            pyg = from_networkx(nx_graph, group_node_attrs=['feature'])
            pyg.catalog = torch.tensor(graph_features['Catalog Code'][index], dtype=torch.long)
            pyg.compiled = torch.tensor(graph_features['compiled'][index], dtype=torch.long)
            pyg.simulated = torch.tensor(graph_features['simulated'][index], dtype=torch.long)
            pyg.J = torch.tensor(graph_features['Utility'][index], dtype=torch.float32)
            pyg.j1 = torch.tensor(graph_features['min J1'][index], dtype=torch.float32)
            pyg.j2 = torch.tensor(graph_features['min J2'][index], dtype=torch.float32)
            pyg.j3 = torch.tensor(graph_features['min J3'][index], dtype=torch.float32)
            pyg.j4 = torch.tensor(graph_features['min J4'][index], dtype=torch.float32)
            pyg.j5 = torch.tensor(graph_features['min J5'][index], dtype=torch.float32)
            if self.pre_filter is not None and not self.pre_filter(pyg):
                continue

            if self.pre_transform is not None:
                data = self.pre_transform(pyg)

            torch.save(pyg, osp.join(self.processed_dir, f'tms_graph_{idx}.pt'))
            idx += 1

    def len(self):
        return len(self.processed_file_names)

    def get(self, idx):
        data = torch.load(f'{self.data_path}/processed/tms_graph_{idx}.pt')
        return data

UPDATE: After looking further into it, after the process has occurred the first time, any time I run

graphs = TMSDataset(root='path')

graphs then equal TMSDataset(32612). What is the purpose of 'get()' then?

[wsad1]

Your code looks ok to me.
Also, get is expected to return a single graph, as described here