core.datasets.base_dataset

Copyright (c) Meta, Inc. and its affiliates. This source code is licensed under the MIT license found in the LICENSE file in the root directory of this source tree.

Attributes

T_co

Exceptions

UnsupportedDatasetError

Inappropriate argument value (of correct type).

Classes

DatasetMetadata
BaseDataset	Base Dataset class for all OCP datasets.
Subset	A pytorch subset that also takes metadata if given.

Functions

create_dataset(→ Subset)

Create a dataset from a config dictionary

Module Contents

core.datasets.base_dataset.T_co

class core.datasets.base_dataset.DatasetMetadata

Bases: NamedTuple

natoms: numpy.typing.ArrayLike | None = None

exception core.datasets.base_dataset.UnsupportedDatasetError

Bases: ValueError

Inappropriate argument value (of correct type).

class core.datasets.base_dataset.BaseDataset(config: dict)

Bases: torch.utils.data.Dataset[T_co]

Base Dataset class for all OCP datasets.

config

paths = []

lin_ref = None

__len__() → int

metadata_hasattr(attr) → bool

property indices

property _metadata: DatasetMetadata

get_metadata(attr, idx)

class core.datasets.base_dataset.Subset(dataset: BaseDataset, indices: collections.abc.Sequence[int], metadata: DatasetMetadata | None = None)

Bases: torch.utils.data.Subset, BaseDataset

A pytorch subset that also takes metadata if given.

metadata

indices

num_samples

config

property _metadata: DatasetMetadata

get_metadata(attr, idx)

core.datasets.base_dataset.create_dataset(config: dict[str, Any], split: str) → Subset

Create a dataset from a config dictionary

Parameters:

• config (dict) – dataset config dictionary

• split (str) – name of split

Returns:

dataset subset class

Return type:

Subset