Class TDataset

code/dataset_benchmark.py:204–257

class TDataset(InMemoryDataset):
    '''

    '''

    def __init__(self, root, name, transform=None, pre_transform=None):
        self.name = name
        assert self.name in ['tfinance', 'tsocial']

        super().__init__(root, transform, pre_transform)
        self.load(self.processed_paths[0])

    @property
    def raw_dir(self):
        return osp.join(self.root, self.name, 'raw')

    @property
    def processed_dir(self):
        return osp.join(self.root, self.name, 'processed')

    @property
    def raw_file_names(self):
        names = [self.name]
        return names

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

    def download(self):
        pass

    def process(self):
        file_path = os.path.join(self.raw_dir, self.raw_file_names[0])
        if not osp.exists(file_path):
            try:
                shutil.copy(f'data/{self.name}', file_path)
            except:
                raise ValueError('source file does not exist!')

        data = load_graphs(file_path)[0][0]
        features = data.ndata['feature']
        labels = data.ndata['label']
        train_mask = data.ndata['train_masks']
        val_mask = data.ndata['val_masks']
        test_mask = data.ndata['test_masks']

        data = Data(x=features, edge_index=torch.vstack(
            data.edges()), y=labels)
        data.tran_mask = train_mask
        data.val_mask = val_mask
        data.test_mask = test_mask
        data = data if self.pre_transform is None else self.pre_transform(data)
        self.save([data], self.processed_paths[0])


class TextDataset(InMemoryDataset):