Skip to content

BaseSyntheticDataset

Methods and Attributes

Bases: Dataset, ABC

Abstract base class for generating 2D synthetic datasets.

Parameters:

Name Type Description Default
n_samples int

The total number of samples to generate.

 required
device Optional[Union[str, device]]

The device to place the tensor on.

 None
dtype dtype

The data type for the output tensor.

 float32
seed Optional[int]

A random seed for reproducibility.

 None
Source code in torchebm/datasets/generators.py 
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
class BaseSyntheticDataset(Dataset, ABC):
    """
    Abstract base class for generating 2D synthetic datasets.

    Args:
        n_samples (int): The total number of samples to generate.
        device (Optional[Union[str, torch.device]]): The device to place the tensor on.
        dtype (torch.dtype): The data type for the output tensor.
        seed (Optional[int]): A random seed for reproducibility.
    """

    def __init__(
        self,
        n_samples: int,
        device: Optional[Union[str, torch.device]] = None,
        dtype: torch.dtype = torch.float32,
        seed: Optional[int] = None,
    ):
        if n_samples <= 0:
            raise ValueError("n_samples must be positive")

        self.n_samples = n_samples
        self.device = device
        self.dtype = dtype
        self.seed = seed
        self.data: Optional[torch.Tensor] = None  # Data will be stored here
        self._generate()  # Generate data upon initialization

    def _seed_generators(self):
        """Sets the random seeds for numpy and torch if a seed is provided."""
        if self.seed is not None:
            np.random.seed(self.seed)
            torch.manual_seed(self.seed)
            # If using CUDA, also seed the CUDA generator
            if torch.cuda.is_available() and (
                isinstance(self.device, torch.device)
                and self.device.type == "cuda"
                or self.device == "cuda"
            ):
                torch.cuda.manual_seed_all(self.seed)  # Seed all GPUs

    @abstractmethod
    def _generate_data(self) -> torch.Tensor:
        """
        Core data generation logic to be implemented by subclasses.
        """
        pass

    def _generate(self):
        """Internal method to handle seeding and call the generation logic."""
        self._seed_generators()
        # Generate data using the subclass implementation
        generated_output = self._generate_data()

        # Ensure it's a tensor and on the correct device/dtype
        if isinstance(generated_output, np.ndarray):
            self.data = _to_tensor(
                generated_output, dtype=self.dtype, device=self.device
            )
        elif isinstance(generated_output, torch.Tensor):
            self.data = generated_output.to(dtype=self.dtype, device=self.device)
        else:
            raise TypeError(
                f"_generate_data must return a NumPy array or PyTorch Tensor, got {type(generated_output)}"
            )

        # Verify batch_shape
        if self.data.shape[0] != self.n_samples:
            warnings.warn(
                f"Generated data has {self.data.shape[0]} samples, but {self.n_samples} were requested. Check generation logic.",
                RuntimeWarning,
            )
            # Optional: adjust self.n_samples or raise error depending on desired strictness
            # self.n_samples = self.data.batch_shape[0]

    def regenerate(self, seed: Optional[int] = None):
        """
        Re-generates the dataset, optionally with a new seed.

        Args:
            seed (Optional[int]): A new random seed. If `None`, the original seed is used.
        """
        if seed is not None:
            self.seed = seed  # Update the seed if a new one is provided
        self._generate()

    def get_data(self) -> torch.Tensor:
        """
        Returns the entire generated dataset as a single tensor.
        """
        if self.data is None:
            # Should not happen if _generate() is called in __init__
            self._generate()
        return self.data

    def __len__(self) -> int:
        """Returns the number of samples in the dataset."""
        return self.n_samples

    def __getitem__(self, idx: int) -> torch.Tensor:
        """
        Returns the sample at the specified index.
        """
        if self.data is None:
            self._generate()  # Ensure data exists

        if not 0 <= idx < self.n_samples:
            raise IndexError(
                f"Index {idx} out of bounds for dataset with size {self.n_samples}"
            )
        return self.data[idx]

    def __repr__(self) -> str:
        """String representation of the dataset object."""
        params = [f"n_samples={self.n_samples}"]
        # Add specific params from subclasses if desired, e.g. by inspecting self.__dict__
        # Or define __repr__ in subclasses
        return f"{self.__class__.__name__}({', '.join(params)}, device={self.device}, dtype={self.dtype})"

n_samples instance-attribute 

n_samples = n_samples

device instance-attribute 

device = device

dtype instance-attribute 

dtype = dtype

seed instance-attribute 

seed = seed

data instance-attribute 

data: Optional[Tensor] = None

_seed_generators 

_seed_generators()

Sets the random seeds for numpy and torch if a seed is provided.

Source code in torchebm/datasets/generators.py 
57
58
59
60
61
62
63
64
65
66
67
68
def _seed_generators(self):
    """Sets the random seeds for numpy and torch if a seed is provided."""
    if self.seed is not None:
        np.random.seed(self.seed)
        torch.manual_seed(self.seed)
        # If using CUDA, also seed the CUDA generator
        if torch.cuda.is_available() and (
            isinstance(self.device, torch.device)
            and self.device.type == "cuda"
            or self.device == "cuda"
        ):
            torch.cuda.manual_seed_all(self.seed)  # Seed all GPUs

_generate_data abstractmethod 

_generate_data() -> torch.Tensor

Core data generation logic to be implemented by subclasses.

Source code in torchebm/datasets/generators.py 
70
71
72
73
74
75
@abstractmethod
def _generate_data(self) -> torch.Tensor:
    """
    Core data generation logic to be implemented by subclasses.
    """
    pass

_generate 

_generate()

Internal method to handle seeding and call the generation logic.

Source code in torchebm/datasets/generators.py 
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
def _generate(self):
    """Internal method to handle seeding and call the generation logic."""
    self._seed_generators()
    # Generate data using the subclass implementation
    generated_output = self._generate_data()

    # Ensure it's a tensor and on the correct device/dtype
    if isinstance(generated_output, np.ndarray):
        self.data = _to_tensor(
            generated_output, dtype=self.dtype, device=self.device
        )
    elif isinstance(generated_output, torch.Tensor):
        self.data = generated_output.to(dtype=self.dtype, device=self.device)
    else:
        raise TypeError(
            f"_generate_data must return a NumPy array or PyTorch Tensor, got {type(generated_output)}"
        )

    # Verify batch_shape
    if self.data.shape[0] != self.n_samples:
        warnings.warn(
            f"Generated data has {self.data.shape[0]} samples, but {self.n_samples} were requested. Check generation logic.",
            RuntimeWarning,
        )

regenerate 

regenerate(seed: Optional[int] = None)

Re-generates the dataset, optionally with a new seed.

Parameters:

Name Type Description Default
seed Optional[int]

A new random seed. If None, the original seed is used.

 None
Source code in torchebm/datasets/generators.py 
104
105
106
107
108
109
110
111
112
113
def regenerate(self, seed: Optional[int] = None):
    """
    Re-generates the dataset, optionally with a new seed.

    Args:
        seed (Optional[int]): A new random seed. If `None`, the original seed is used.
    """
    if seed is not None:
        self.seed = seed  # Update the seed if a new one is provided
    self._generate()

get_data 

get_data() -> torch.Tensor

Returns the entire generated dataset as a single tensor.

Source code in torchebm/datasets/generators.py 
115
116
117
118
119
120
121
122
def get_data(self) -> torch.Tensor:
    """
    Returns the entire generated dataset as a single tensor.
    """
    if self.data is None:
        # Should not happen if _generate() is called in __init__
        self._generate()
    return self.data