Skip to content

dataset

Dataset

Bases: InMemoryDataset

In-memory dataset that loads a pre-processed, collated set of StructData samples from disk.

Source code in src/torch_structure/data/dataset.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
class Dataset(InMemoryDataset):
    """In-memory dataset that loads a pre-processed, collated set of [StructData][torch_structure.data.data.StructData] samples from disk."""

    def __init__(self, root, load=True, transform=None):
        super().__init__(root, transform=transform)
        if load and os.path.exists(self.processed_paths[0]):
            self.load(self.processed_paths[0])
        else:
            self.data, self.slices = None, None

    @property
    def processed_file_names(self):
        """Filename of the collated dataset file, relative to ``root/processed``."""
        return ["data.pt"]

processed_file_names property

Filename of the collated dataset file, relative to root/processed.

LegacyDataset

Bases: InMemoryDataset

In-memory dataset that stores and returns a list of pre-loaded StructData samples as-is.

Source code in src/torch_structure/data/dataset.py 
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
class LegacyDataset(InMemoryDataset):
    """In-memory dataset that stores and returns a list of pre-loaded [StructData][torch_structure.data.data.StructData] samples as-is."""

    def __init__(self, data_list):
        super().__init__(None)
        self.data_list = data_list
        self.data, self.slices = self.collate(self.data_list)  # double check this

    def len(self):
        """Return the number of samples in the dataset."""
        return len(self.data_list)

    def get(self, idx):
        """Return the sample at ``idx``."""
        return self.data_list[idx]  # .data

get(idx)

Return the sample at idx.

Source code in src/torch_structure/data/dataset.py 
18
19
20
def get(self, idx):
    """Return the sample at ``idx``."""
    return self.data_list[idx]  # .data

len()

Return the number of samples in the dataset.

Source code in src/torch_structure/data/dataset.py 
14
15
16
def len(self):
    """Return the number of samples in the dataset."""
    return len(self.data_list)

save(data_list, root, include_metadata=False)

Collate a list of StructData samples and save them as a Dataset under root.

Parameters:

Name Type Description Default
data_list list[StructData]

samples to save.

 required
root str

dataset root directory; the collated file is written to root/processed/data.pt.

 required
include_metadata bool

if False (default), strips each sample's metadata attribute before saving.

 False
Source code in src/torch_structure/data/dataset.py 
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
def save(data_list, root, include_metadata=False):
    """Collate a list of [StructData][torch_structure.data.data.StructData] samples and save them as a [Dataset][torch_structure.data.dataset.Dataset] under ``root``.

    Args:
        data_list (list[StructData]): samples to save.
        root (str): dataset root directory; the collated file is written to
            ``root/processed/data.pt``.
        include_metadata (bool): if ``False`` (default), strips each
            sample's ``metadata`` attribute before saving.
    """
    dataset = Dataset(root, load=False)
    if not include_metadata:
        for data in data_list:
            del data.metadata
    dataset.save(data_list, dataset.processed_paths[0])