How much RAM is needed to load the TPUGraphs dataset?

[Jackmin801]

Hello authors!

I tried to run the script using the given command

python main_tpugraphs.py --cfg configs/tpugraphs.yaml

However, the script seems to get stuck trying to load the data, consuming all the RAM on my machine (240GB) and then being killed by the kernel.

Does this make sense? The dataset is only 5.8GB on disk, albeit zipped. But it expanding to more than 200GB seems quite off.

[Jackmin801]

Also, which version of torch-geometric should be used to run the script?
I am currently using torch-geometric==2.3.1 but it does not have the configs cfg.train.mode and cfg.device.
I had to add these flags to get the script to get through the config parsing portion

[gabrielgramm]

Hey Jackmin
Have you managed to run the whole code with torch-geometric==2.3.1, if yes which installation settings have you used for the pyg-lib?
The one stated in the GPS description dont work for me.
(pip install pyg-lib -f https://data.pyg.org/whl/torch-1.13.0+cu117.html)

[SongWang2017]

Hey Jackmin Have you managed to run the whole code with torch-geometric==2.3.1, if yes which installation settings have you used for the pyg-lib? The one stated in the GPS description dont work for me. (pip install pyg-lib -f https://data.pyg.org/whl/torch-1.13.0+cu117.html)

Hi Gabriel:

I think there're some conflicts between the latest version of torch-geometric and torch under version 1.13+cuda, fortunately I found that torch-2.0.0-cu117 works

!rm -rf /opt/conda/lib/python3.10/site-packages/torch*

!pip install torch torch_geometric pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.0.0+cu117.html &> /dev/null

Hoping this would help

[SongWang2017]

Also, which version of torch-geometric should be used to run the script? I am currently using torch-geometric==2.3.1 but it does not have the configs cfg.train.mode and cfg.device. I had to add these flags to get the script to get through the config parsing portion

Hi Jackmin:

Seems it is another version problem with set_cfg, the latest version of set_cfg from torch_geometric.graphgym.config just doesn't allow accepting new cfg attributes, I'll issue another pull request and hope that would help.

[SongWang2017]

Hello authors!

I tried to run the script using the given command

python main_tpugraphs.py --cfg configs/tpugraphs.yaml

However, the script seems to get stuck trying to load the data, consuming all the RAM on my machine (240GB) and then being killed by the kernel.

Does this make sense? The dataset is only 5.8GB on disk, albeit zipped. But it expanding to more than 200GB seems quite off.

Hi Jackmin:

This comes as no surprise, the dataset is only about nodes/edges and features(mainly numpy arrays) so its disk usage is quite small, however the data processing pipeline construct graphs that cost much more memory, the difference is just like two numpy array to store Vertices and Edges versus a full 2D array to describe the adjacent information between them, and the gap between them will grow larger as the number of nodes gets larger

If you save the graphs(Data class) one by one, you'd find that even a single graph would consume gigabytes of disk space! I experimentally saved all fully constructed graphs(objects of class Data) of all source/search combinations, a disk space of about 400gb is required.

The TPUGraphs dataset tries to collect all graphs (in memory) and save them into a single pt file, so the system needs hundreds of Gigabytes to store them.

To circumvent this, I think you could just reconstruct the graph one by one as shown by the following code

def get(self, idx):        
        filename = self.processed_file_names[idx]
        np_file = dict(np.load(filename))
        if "edge_index" not in np_file:
            print('error in', filename)
        edge_index = torch.tensor(np_file["edge_index"].T)
        runtime = torch.tensor(np_file["config_runtime"])
        op = torch.tensor(np_file["node_feat"])
        op_code = torch.tensor(np_file["node_opcode"])
        config_feats = torch.tensor(np_file["node_config_feat"])
        config_feats = config_feats.view(-1, config_feats.shape[-1])
        config_idx = torch.tensor(np_file["node_config_ids"])
        num_config = torch.tensor(np_file["node_config_feat"].shape[0])
        num_config_idx = torch.tensor(np_file["node_config_feat"].shape[1])
        num_nodes = torch.tensor(np_file["node_feat"].shape[0])
        num_parts = num_nodes // self.thres + 1
        interval = num_nodes // num_parts
        partptr = torch.arange(0, num_nodes, interval+1)
        if partptr[-1] != num_nodes:
            partptr = torch.cat([partptr, torch.tensor([num_nodes])])
        data = Data(edge_index=edge_index, op_feats=op, op_code=op_code, config_feats=config_feats, config_idx=config_idx,
                    num_config=num_config, num_config_idx=num_config_idx, y=runtime, num_nodes=num_nodes, partptr=partptr, partition_idx=self.parts_cnt)
        self.parts_cnt += num_parts * num_config
        
        op_feats_mean = torch.mean(data.op_feats, dim=0, keepdim=True)
        op_feats_std = torch.std(data.op_feats, dim=0, keepdim=True)
        op_feats_std[op_feats_std < 1e-6] = 1
        data.op_feats = (data.op_feats - op_feats_mean) / op_feats_std
                
        return data

[gabrielgramm]

Thanks for the reply. But where has this ‘get’ function to be added?