Have a pytorch class produce 2 separate datasets - pytorch

Question

Does anyone have any suggestions for creating 2 datasets at once? When I run dataset = ImbCircuitDataset('test'), I want it to save 2 separate sets, even if I have to go all the way back to self.data and split it there, I just want the class to produce 2 separate datasets.
Basically run set1, set2 = ImbCircuitDataset('test')
I started doing it at the bottom of the def process(self) but I have a feeling its wrong. I am thinking in def processed_file_names(self) I need to return 'set1.pt', 'set2.pt'.
from scipy.io import loadmat
import pandas as pd
from torch_geometric.data import Data, InMemoryDataset
import torch
import neptune.new as neptune
import numpy as np
import torch_geometric
from tqdm import tqdm
from typing import Tuple, Union
from torch import Tensor
from collections.abc import Sequence
from torch_geometric.utils import from_networkx
import networkx as nx
import matplotlib.pyplot as plt
IndexType = Union[slice, Tensor, np.ndarray, Sequence]
print(f"Torch version: {torch.__version__}")
print(f"Cuda available: {torch.cuda.is_available()}")
print(f"Torch geometric version: {torch_geometric.__version__}")
class ImbCircuitDataset(InMemoryDataset):
def __init__(self, root, transform=None, pre_transform=None,
pre_filter=None):
super().__init__(root, transform, pre_transform, pre_filter)
self.data, self.slices = torch.load(self.processed_paths[0])
#property
def raw_file_names(self):
return 'shuffled_data.mat'
#property
def processed_file_names(self):
return 'data.pt'
def download(self):
pass
def process(self):
self.raw_data = loadmat(self.raw_paths[0], squeeze_me=True)
self.data = pd.DataFrame(self.raw_data['Graphs'])
self.data = self.data.sort_values('Labels', ascending=True, ignore_index=True)
for i in range(len(self.data)):
self.data['Ln'][i] = dict(enumerate(self.data['Ln'][i]))
data_list = []
for index, cir in tqdm(self.data.iterrows(), total=self.data.shape[0]):
nxg = nx.Graph(self.data['A'][index])
nx.set_node_attributes(nxg, self.data['Ln'][index], 'component')
pt_graph = self._get_graph_object(nxg)
pt_graph.x = self._get_node_features(nxg, self.data['Ln'][index])
pt_graph.performance = torch.tensor(self.data['Labels'][index], dtype=torch.float)
data_list.append(pt_graph)
if self.pre_filter is not None:
data_list = [d for d in data_list if self.pre_filter(d)]
if self.pre_transform is not None:
data_list = [self.pre_transform(d) for d in data_list]
split = 0.5
N = len(data_list)
set_1 = data_list[:N*split]
set_2 = data_list[N*split:]
data1, slices1 = self.collate(set_1)
data2, slices2 = self.collate(set_2)
torch.save((data1, slices1), self.processed_paths[0])
torch.save((data2,slices2), self.processed_paths[0])
def _get_node_features(self, nxgraph, node_labels):
betweenness = list(nx.betweenness_centrality(nxgraph).values())
eigenvector = list(nx.eigenvector_centrality(
nxgraph, max_iter=600).values())
mapping_dict = {'C': 0, 'G': 1, 'I': 2, 'O': 3, 'P': 4, 'R': 5}
component_labels = []
for value in node_labels.values():
if value in mapping_dict:
component_labels.append(mapping_dict[value])
all_feats = [component_labels, betweenness, eigenvector]
all_feats = np.asarray(all_feats).transpose()
return torch.tensor(all_feats, dtype=torch.float)
def _get_graph_object(self, nx_graph):
nxg = from_networkx(nx_graph)
return nxg
#property
def num_node_features(self) -> int:
return 3
#property
def num_classes(self) -> int:
return 2