From cf0478a489cfbe450bfd7433a05cd80446048e82 Mon Sep 17 00:00:00 2001 From: amansouribigvand Date: Thu, 27 Nov 2025 21:43:21 +0000 Subject: [PATCH 1/2] Add RBM --- .gitignore | 8 +- .../plugins/torch/models/boltzmann_machine.py | 276 +++++++++++++++++- examples/rbm_image_generation.py | 214 ++++++++++++++ tests/test_boltzmann_machine.py | 272 +++++++++++++++++ 4 files changed, 768 insertions(+), 2 deletions(-) create mode 100644 examples/rbm_image_generation.py diff --git a/.gitignore b/.gitignore index 39c649a..bac9255 100644 --- a/.gitignore +++ b/.gitignore @@ -66,4 +66,10 @@ venv.bak/ dmypy.json # aim -*.aim* \ No newline at end of file +*.aim* + +# Datasets +data/* + +# Generated images +samples/* \ No newline at end of file diff --git a/dwave/plugins/torch/models/boltzmann_machine.py b/dwave/plugins/torch/models/boltzmann_machine.py index 8bd0350..e06f518 100644 --- a/dwave/plugins/torch/models/boltzmann_machine.py +++ b/dwave/plugins/torch/models/boltzmann_machine.py @@ -43,7 +43,7 @@ spread = AggregatedSamples.spread -__all__ = ["GraphRestrictedBoltzmannMachine"] +__all__ = ["GraphRestrictedBoltzmannMachine", "RestrictedBoltzmannMachine"] class GraphRestrictedBoltzmannMachine(torch.nn.Module): @@ -662,3 +662,277 @@ def estimate_beta(self, spins: torch.Tensor) -> float: bqm = BinaryQuadraticModel.from_ising(*self.to_ising(1)) beta = 1 / mple(bqm, (spins.detach().cpu().numpy(), self._nodes))[0] return beta + +class RestrictedBoltzmannMachine(torch.nn.Module): + """A Restricted Boltzmann Machine (RBM) model. + + This class defines the parameterization and inference of a binary RBM. + Training is performed using Persistent Contrastive Divergence (PCD). + + Args: + n_visible (int): Number of visible units. + n_hidden (int): Number of hidden units. + """ + + def __init__( + self, + n_visible: int, + n_hidden: int, + ) -> None: + super().__init__() + + # Model hyperparameters + self._n_visible = n_visible + self._n_hidden = n_hidden + + # Initialize model parameters + # initialize weights + self._weights = torch.nn.Parameter( + 0.1 * torch.randn(n_visible, n_hidden) + ) + # initialize visible units biases. + self._visible_biases = torch.nn.Parameter( + 0.5 * torch.ones(n_visible) + ) + # initialize hidden units biases. + self._hidden_biases = torch.nn.Parameter( + 0.5 * torch.ones(n_hidden) + ) + + # Stores the last visible states to initialize the Markov chain in Persistent Contrastive Divergence (PCD) + self.register_buffer("_previous_visible_values", None) + + # Initialize momenta tensors for momentum-based updates (all start at 0) + self.register_buffer("_weight_momenta", torch.zeros(n_visible, n_hidden)) + self.register_buffer("_visible_bias_momenta", torch.zeros(n_visible)) + self.register_buffer("_hidden_bias_momenta", torch.zeros(n_hidden)) + + @property + def n_visible(self) -> int: + """Number of visible units.""" + return self._n_visible + + @property + def n_hidden(self) -> int: + """Number of hidden units.""" + return self._n_hidden + + @property + def weights(self) -> torch.Tensor: + """Weights of the RBM.""" + return self._weights + + @property + def visible_biases(self) -> torch.Tensor: + """Visible biases of the RBM.""" + return self._visible_biases + + @property + def hidden_biases(self) -> torch.Tensor: + """Hidden biases of the RBM.""" + return self._hidden_biases + + @property + def previous_visible_values(self) -> torch.Tensor: + """Previous visible values used in Persistent Contrastive Divergence (PCD).""" + return self._previous_visible_values + + @property + def weight_momenta(self) -> torch.Tensor: + """Weight momenta of the RBM.""" + return self._weight_momenta + + @property + def visible_bias_momenta(self) -> torch.Tensor: + """Visible bias momenta of the RBM.""" + return self._visible_bias_momenta + + @property + def hidden_bias_momenta(self) -> torch.Tensor: + """Hidden bias momenta of the RBM.""" + return self._hidden_bias_momenta + + + def _sample_hidden(self, visible: torch.Tensor) -> torch.Tensor: + """Sample from the distribution P(h|v). + + Args: + visible (torch.Tensor): Tensor of shape (batch_size, n_visible) + representing the states of visible units. + + Returns: + torch.Tensor: Binary tensor of shape (batch_size, n_hidden) representing + sampled hidden units. + """ + hidden_probs = torch.sigmoid(self._hidden_biases + visible @ self._weights) + return torch.bernoulli(hidden_probs) + + def _sample_visible(self, hidden: torch.Tensor) -> torch.Tensor: + """Sample from the distribution P(v|h). + + Args: + hidden (torch.Tensor): Tensor of shape (batch_size, n_hidden) + representing the states of hidden units. + Returns: + torch.Tensor: Binary tensor of shape (batch_size, n_visible) representing + sampled visible units. + """ + visible_probs = torch.sigmoid(self._visible_biases + hidden @ self._weights.t()) + return torch.bernoulli(visible_probs) + + def generate_sample( + self, + batch_size: int, + gibbs_steps: int, + start_visible: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Generate a sample of visible and hidden units using gibbs sampling. + + Args: + batch_size (int): Number of samples to generate. + gibbs_steps (int): Number of Gibbs sampling steps to perform. + start_visible (torch.Tensor | None, optional): Initial visible states to + start the Gibbs chain (shape: [batch_size, n_visible]). If None, + a random Gaussian initialization is used. + + + Returns: + tuple[torch.Tensor, torch.Tensor]: A tuple of (visible, hidden) from the last Gibbs step: + - visible: (batch_size, n_visible) + - hidden: (batch_size, n_hidden) + """ + if start_visible is None: + visible_values = torch.randn( + batch_size, self.n_visible, device=self._weights.device + ) + else: + visible_values = start_visible + + hidden_values = None + + for _ in range(gibbs_steps): + hidden_values = self._sample_hidden(visible_values) + visible_values = self._sample_visible(hidden_values) + + return visible_values, hidden_values + + def _contrastive_divergence( + self, + batch: torch.Tensor, + epoch: int, + n_gibbs_steps: int, + learning_rate: float, + momentum_coefficient: float, + weight_decay: float, + n_epochs: int, + ) -> torch.Tensor: + """ + Perform one step of Contrastive Divergence (CD-k) with momentum and weight decay. + Uses Persistent Contrastive Divergence (PCD) by maintaining the last visible states + for Gibbs sampling across batches. + + Args: + batch (torch.Tensor): A batch of input data of shape (batch_size, n_visible). + epoch (int): Current training epoch. + n_gibbs_steps (int): Number of Gibbs sampling steps per epoch. + learning_rate (float): Base learning rate for parameter updates. + momentum_coefficient (float): Momentum coefficient for parameter updates. + weight_decay (float): weight decay (L2 regularization) coefficient for weights. + n_epochs (int): Number of training epochs. + + Returns: + torch.Tensor: The reconstruction error (L1 norm) for the batch. + + """ + + # Positive phase (data-driven) + hidden_probs = torch.sigmoid(self._hidden_biases + batch @ self._weights) + + weight_grads = torch.matmul(batch.t(), hidden_probs) + visible_bias_grads = batch + hidden_bias_grads = hidden_probs + + batch_size = batch.size(0) + + # Initialize previous visible states for Persistent CD + if self._previous_visible_values == None: + self._previous_visible_values = torch.randn_like( + batch, device=self._weights.device + ) + + # Negative phase (model-driven) + # Sample from the model using gibbs sampling + visible_values, hidden_values = self.generate_sample( + batch_size, n_gibbs_steps, self._previous_visible_values + ) + + visible_values = visible_values.detach() + hidden_values = hidden_values.detach() + # Store samples to initialize the next Markov chain with (PCD) + self._previous_visible_values = visible_values + + # Compute the gradients for negative phase + weight_grads -= torch.matmul(visible_values.t(), hidden_values) + + visible_bias_grads -= visible_values + hidden_bias_grads -= hidden_values + + # Average across the batch + weight_grads /= batch_size + visible_bias_grads /= batch_size + hidden_bias_grads /= batch_size + + # Compute decayed learning rate + decayed_learning_rate = learning_rate - (learning_rate / n_epochs * epoch) + + # Update momenta + self._weight_momenta = self._weight_momenta * momentum_coefficient + decayed_learning_rate * weight_grads + self._visible_bias_momenta = self._visible_bias_momenta * momentum_coefficient + decayed_learning_rate * torch.sum( + visible_bias_grads, dim=0 + ) + self._hidden_bias_momenta = self._hidden_bias_momenta * momentum_coefficient + decayed_learning_rate * torch.sum( + hidden_bias_grads, dim=0 + ) + + with torch.no_grad(): + # Update parameters + self._weights += self._weight_momenta + self._visible_biases += self._visible_bias_momenta + self._hidden_biases += self._hidden_bias_momenta + + # Apply weight decay + self._weights -= decayed_learning_rate * self._weights * weight_decay + + # Compute reconstruction error (L1 norm) + reconstruction = self._sample_visible(self._sample_hidden(batch)) + reconstruction = reconstruction.detach() + error = torch.sum(torch.abs(batch - reconstruction)) + + return error + + def forward(self, visible: torch.Tensor) -> torch.Tensor: + """ + Computes the RBM free energy of a batch of visible units averaged over the batch. + + The free energy F(visible) for a visible vector visible is: + + F(visible) = - visible · visible_biases + - sum_{j=1}^{n_hidden} log(1 + exp(hidden_biases[j] + (visible · weights)_j)) + + Args: + visible (torch.Tensor): Tensor of shape (batch_size, n_visible) representing the visible layer. + + Returns: + torch.Tensor: Scalar tensor representing the **average free energy** over the batch. + """ + + v_term = (visible * self._visible_biases).sum(dim=1) + + hidden_pre_activation = visible @ self._weights + self._hidden_biases + + h_term = torch.sum(torch.nn.functional.softplus(hidden_pre_activation), dim=1) + + free_energy_per_sample = -v_term - h_term + + # average over batch + return free_energy_per_sample.mean() diff --git a/examples/rbm_image_generation.py b/examples/rbm_image_generation.py new file mode 100644 index 0000000..bd5f14a --- /dev/null +++ b/examples/rbm_image_generation.py @@ -0,0 +1,214 @@ +# Copyright 2025 D-Wave +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import torch +from torch.utils.data import DataLoader +from dwave.plugins.torch.models.boltzmann_machine import ( + RestrictedBoltzmannMachine as RBM, +) +from torchvision import transforms, datasets +import matplotlib.pyplot as plt + + +def load_binarized_mnist(dataset_path: str = "data") -> datasets.MNIST: + """ + Load the MNIST dataset and binarize it (pixels >= 0.5 become 1, else 0). + + Args: + dataset_path (str): Path to download/store the MNIST dataset. Defaults to "data". + + Returns: + datasets.MNIST: Binarized MNIST training dataset. + """ + transform = transforms.Compose( + [transforms.ToTensor(), transforms.Lambda(lambda x: (x >= 0.5).float())] + ) + + train_dataset = datasets.MNIST( + root=dataset_path, train=True, transform=transform, download=True + ) + return train_dataset + + +def train_loop( + train_loader: DataLoader, + rbm: RBM, + n_epochs: int, + n_gibbs_steps: int, + learning_rate: float, + momentum: float, + weight_decay: float, +) -> None: + """ + Train the RBM using contrastive divergence with momentum and weight decay. + + Args: + train_loader (DataLoader): PyTorch DataLoader for training data. + rbm (RBM): Restricted Boltzmann Machine instance. + n_epochs (int): Number of training epochs. + n_gibbs_steps (int): Number of Gibbs sampling steps per CD update. + learning_rate (float): Base learning rate. + momentum (float): Momentum coefficient for parameter updates. + weight_decay (float): Weight decay (L2 regularization) coefficient. + """ + device = rbm._weights.device + for epoch in range(n_epochs): + total_error = 0 + num_examples = 0 + for batch, _ in train_loader: + # flatten input data + batch = batch.reshape(batch.size(0), rbm.n_visible).to(device) + + # Perform one step of contrastive divergence and accumulate error + error = rbm._contrastive_divergence( + batch, + epoch, + n_gibbs_steps, + learning_rate, + momentum, + weight_decay, + n_epochs, + ) + total_error += error + num_examples += batch.size(0) + average_error = total_error / num_examples # Average reconstruction error + print( + f"Epoch {epoch + 1}/{n_epochs} - Avg reconstruction error: {average_error:.4f}" + ) + + +def generate_and_save_images( + rbm: RBM, + rows: int = 8, + columns: int = 8, + steps: int = 1000, + output_dir: str = "samples", + output_filename: str = "generated_images.png", +) -> None: + """ + Generate samples from a trained RBM and save them as a grid of images. + + Args: + rbm (RBM): Trained RBM instance. + rows (int): Number of rows in the output image grid. Defaults to 8. + columns (int): Number of columns in the output image grid. Defaults to 8. + steps (int): Number of Gibbs sampling steps for generation. Defaults to 1000. + output_dir (str): Directory to save the generated images. Defaults to "samples". + output_filename (str): File name for saving the generated image grid. Defaults to "generated_images.png". + """ + os.makedirs(output_dir, exist_ok=True) + output_path = os.path.join(output_dir, output_filename) + + num_images = rows * columns + + # Generate batch of images + samples, _ = rbm.generate_sample(num_images, gibbs_steps=steps) + samples = samples.view(num_images, 28, 28).detach().cpu().numpy() + + # Plot grid of images + fig, axs = plt.subplots(rows, columns, figsize=(columns, rows)) + + idx = 0 + for r in range(rows): + for c in range(columns): + axs[r, c].imshow(samples[idx], cmap="gray") + axs[r, c].axis("off") + idx += 1 + + fig.suptitle("Generated images from RBM trained on MNIST", fontsize=18) + plt.tight_layout(rect=[0, 0, 1, 0.96]) + plt.savefig(output_path, bbox_inches="tight", pad_inches=0.1) + plt.show() + print(f"Generated {num_images} samples in {output_dir}/{output_filename}") + + +def train_rbm( + n_visible: int = 784, + n_hidden: int = 500, + n_gibbs_steps: int = 10, + learning_rate: float = 1e-3, + momentum: float = 0.5, + weight_decay: float = 1e-7, + n_epochs: int = 20, + batch_size: int = 32, + dataset_path: str = "data", + output_dir: str = "samples", + output_filename: str = "generated_images.png", +) -> None: + """Train an RBM on MNIST and generate sample images. + + Args: + n_visible (int, optional): Number of visible units. Defaults to 784. + n_hidden (int, optional): Number of hidden units. Defaults to 500. + n_gibbs_steps (int, optional): Number of Gibbs sampling steps per CD update. Defaults to 10. + learning_rate (float, optional): Base learning rate for CD updates. Defaults to 1e-3. + momentum (float, optional): Momentum coefficient for CD updates. Defaults to 0.5. + weight_decay (float, optional): Weight decay (L2 regularization) coefficient. Defaults to 1e-7. + n_epochs (int, optional): Number of training epochs. Defaults to 20. + batch_size (int, optional): Batch size for training. Defaults to 32. + dataset_path (str, optional): Path to download/store the MNIST dataset. Defaults to "data". + output_dir (str, optional): Directory to save the generated images. Defaults to "samples". + output_filename (str, optional): File name for saving the generated image grid. Defaults to "generated_images.png". + """ + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + print(f"Using {device} device") + + # Load MNIST data + print("Loading MNIST dataset...") + train_dataset = load_binarized_mnist(dataset_path) + + # Create data loader + train_loader = torch.utils.data.DataLoader( + train_dataset, batch_size=batch_size, shuffle=True, drop_last=True + ) + + # Initialize RBM + rbm = RBM(n_visible, n_hidden).to(device) + + # Train RBM + train_loop( + train_loader, + rbm, + n_epochs=n_epochs, + n_gibbs_steps=n_gibbs_steps, + learning_rate=learning_rate, + momentum=momentum, + weight_decay=weight_decay, + ) + + # Generate and save samples + generate_and_save_images( + rbm, + rows=8, + columns=8, + steps=1000, + output_dir=output_dir, + output_filename=output_filename, + ) + + +if __name__ == "__main__": + # Run an example of fitting a Restricted Boltzmann Machine to the MNIST dataset + train_rbm( + n_visible=784, + n_hidden=500, + n_gibbs_steps=10, + learning_rate=1e-3, + momentum=0.5, + weight_decay=1e-7, + n_epochs=50, + batch_size=64, + ) diff --git a/tests/test_boltzmann_machine.py b/tests/test_boltzmann_machine.py index 069cadf..2e4b297 100644 --- a/tests/test_boltzmann_machine.py +++ b/tests/test_boltzmann_machine.py @@ -21,6 +21,7 @@ from dwave.plugins.torch.models.boltzmann_machine import GraphRestrictedBoltzmannMachine as GRBM from dwave.system.temperatures import maximum_pseudolikelihood_temperature as mple +from dwave.plugins.torch.models.boltzmann_machine import RestrictedBoltzmannMachine as RBM class TestGraphRestrictedBoltzmannMachine(unittest.TestCase): def setUp(self) -> None: @@ -407,6 +408,277 @@ def test_quasi_objective_gradient_hidden_units(self): # the sufficient statistics of the average spins. torch.testing.assert_close(grad, grad_auto) +class TestRBM(unittest.TestCase): + def setUp(self): + # Small RBM for testing + self.rbm = RBM(n_visible=4, n_hidden=3) + + # Common input data for CD tests + self.batch = torch.tensor([[1.0, 0.0, 1.0, 1.0], [0.0, 1.0, 0.0, 1.0]]) + + # Shared CD kwargs + self.cd_kwargs = dict( + epoch=0, + n_gibbs_steps=1, + learning_rate=0.1, + momentum_coefficient=0.5, + weight_decay=0.0, + n_epochs=10, + ) + + def test_sample_hidden_shape(self): + visible = torch.randint(0, 2, (5, self.rbm.n_visible)).float() + hidden = self.rbm._sample_hidden(visible) + # Ensure shape is correct + self.assertEqual(hidden.shape, (5, self.rbm.n_hidden)) + + def test_sample_hidden_binary(self): + visible = torch.randint(0, 2, (5, self.rbm.n_visible)).float() + hidden = self.rbm._sample_hidden(visible) + + # Ensure output is binary + self.assertTrue(torch.all((hidden == 0) | (hidden == 1))) + + @parameterized.expand( + [ + ("all_ones", 1000.0, 1), + ("all_zeroes", -1000.0, 0), + ] + ) + def test_sample_hidden_saturation(self, name, bias_value, expected_value): + """ + Test that _sample_hidden saturates correctly when the hidden biases + are set to very large positive or negative values. + + If hidden_bias[j] → +∞ + sigmoid(hidden_bias + visible @ weights) → 1 + bernoulli(1) → always 1 + + If hidden_bias[j] → -∞ + sigmoid(hidden_bias + visible @ weights) → 0 + bernoulli(0) → always 0 + """ + + # Set all hidden biases to an extreme constant + with torch.no_grad(): + self.rbm._hidden_biases.fill_(bias_value) + + # The visible input does not matter in saturation conditions + visible = torch.zeros(5, self.rbm.n_visible) + + # Sample hidden units + hidden = self.rbm._sample_hidden(visible) + + # Assert that all hidden units match the expected saturated value + self.assertTrue(torch.all(hidden == expected_value)) + + def test_sample_visible_shape(self): + hidden = torch.randint(0, 2, (5, self.rbm.n_hidden)).float() + visible = self.rbm._sample_visible(hidden) + + # Ensure shape is correct + self.assertEqual(visible.shape, (5, self.rbm.n_visible)) + + def test_sample_visible_binary(self): + hidden = torch.randint(0, 2, (5, self.rbm.n_hidden)).float() + visible = self.rbm._sample_visible(hidden) + + # Ensure output is binary + self.assertTrue(torch.all((visible == 0) | (visible == 1)).item()) + + @parameterized.expand( + [ + ("all_ones", 1000.0, 1), + ("all_zeroes", -1000.0, 0), + ] + ) + def test_sample_visible_saturation(self, name, bias_value, expected_value): + """ + Test that _sample_visible saturates correctly when the visible biases + are set to very large positive or negative values. + + If visible_bias → +∞: sigmoid → 1 → bernoulli(1) → always 1 + If visible_bias → -∞: sigmoid → 0 → bernoulli(0) → always 0 + """ + + # Large positive/negative bias makes sigmoid output deterministic + with torch.no_grad(): + self.rbm._visible_biases.fill_(bias_value) + + # Hidden input doesn't matter when biases dominate + hidden = torch.zeros(5, self.rbm.n_hidden) + + visible = self.rbm._sample_visible(hidden) + + self.assertTrue(torch.all(visible == expected_value).item()) + + def test_generate_sample_shape(self): + batch_size = 4 + visible, hidden = self.rbm.generate_sample(batch_size, gibbs_steps=2) + + # Ensure shapes are correct + self.assertEqual(visible.shape, (batch_size, self.rbm.n_visible)) + self.assertEqual(hidden.shape, (batch_size, self.rbm.n_hidden)) + + def test_generate_sample_binary(self): + batch_size = 4 + visible, hidden = self.rbm.generate_sample(batch_size, gibbs_steps=2) + + # Ensure outputs are binary + self.assertTrue(torch.all((visible == 0) | (visible == 1)).item()) + self.assertTrue(torch.all((hidden == 0) | (hidden == 1)).item()) + + def test_generate_sample_initial_hidden_start_visible(self): + """ + When providing a start_visible tensor and using gibbs_steps=1, the first + hidden sample should match _sample_hidden(start_visible). This ensures + the RBM correctly uses the provided initial visible state. + """ + batch_size = 4 + start = torch.zeros(batch_size, self.rbm.n_visible) + + torch.manual_seed(42) + _, hidden = self.rbm.generate_sample( + batch_size, gibbs_steps=1, start_visible=start + ) + + torch.manual_seed(42) + expected_hidden = self.rbm._sample_hidden(start) + + # First hidden sample should match _sample_hidden(start) + self.assertTrue(torch.equal(hidden, expected_hidden)) + + def test_generate_sample_more_gibbs_steps_changes_output(self): + batch_size = 4 + + torch.manual_seed(42) + v1, _ = self.rbm.generate_sample(batch_size, gibbs_steps=1) + + torch.manual_seed(42) + v4, _ = self.rbm.generate_sample(batch_size, gibbs_steps=4) + + # With same seed but longer chain, results should differ + self.assertFalse(torch.allclose(v1, v4)) + + def test_generate_sample_deterministic_small_rbm(self): + """ + Test generate_sample on a tiny deterministic RBM with known parameters. + This test uses a 3x3 RBM with manually set weights, hidden biases, and + visible biases to ensure that the outputs are fully predictable. + """ + rbm = RBM(n_visible=3, n_hidden=3) + + # Set deterministic weights and biases + rbm._weights.data = torch.tensor( + [[0.2, -0.1, 0.0], [0.1, 0.3, -0.2], [-0.2, 0.1, 0.2]] + ) + rbm._hidden_biases.data = torch.tensor([0.1, -0.2, 0.0]) + rbm._visible_biases.data = torch.tensor([0.5, 0.0, 0.3]) + + start_visible = torch.zeros(1, 3) + torch.manual_seed(42) + visible, hidden = rbm.generate_sample( + batch_size=1, gibbs_steps=1, start_visible=start_visible + ) + + # Expected values computed manually + expected_hidden = torch.tensor([[0.0, 0.0, 1.0]]) + expected_visible = torch.tensor([[0.0, 1.0, 1.0]]) + + # Check equality + self.assertTrue(torch.equal(hidden, expected_hidden)) + self.assertTrue(torch.equal(visible, expected_visible)) + + def test_contrastive_divergence_returns_nonnegative_tensor(self): + """CD returns a tensor of type torch.Tensor and L1 error >= 0""" + error = self.rbm._contrastive_divergence( + batch = self.batch, + epoch=0, + n_gibbs_steps=1, + learning_rate=0.01, + momentum_coefficient=0.5, + weight_decay=0.0, + n_epochs=1, + ) + + self.assertIsInstance(error, torch.Tensor) + self.assertGreaterEqual(error.item(), 0.0) + + @parameterized.expand( + [ + ("weights", "_weights"), + ("visible_biases", "_visible_biases"), + ("hidden_biases", "_hidden_biases"), + ] + ) + def test_cd_parameter_updates(self, name, attr): + """Test that the _contrastive_divergence function updates model parameters.""" + torch.manual_seed(42) + before = getattr(self.rbm, attr).clone() + + _ = self.rbm._contrastive_divergence(self.batch, **self.cd_kwargs) + + after = getattr(self.rbm, attr) + + # Check that parameters were updated + self.assertFalse( + torch.allclose(before, after), f"{attr} should update during CD" + ) + + def test_forward_scalar_output(self): + """Forward should return a scalar tensor.""" + batch = torch.randn(5, self.rbm.n_visible) + out = self.rbm.forward(batch) + self.assertIsInstance(out, torch.Tensor) + self.assertEqual(out.ndim, 0) + + def test_forward_zero_weights_biases(self): + """ + Check free energy when all weights and biases are zero. + Analytic test: all weights & biases = 0 + Free energy becomes: + F(v) = - sum_j softplus(0) = -n_hidden * log(2) + """ + with torch.no_grad(): + self.rbm._weights[:] = 0 + self.rbm._visible_biases[:] = 0 + self.rbm._hidden_biases[:] = 0 + v = torch.tensor([[1.0, 0.0, 1.0, 1.0]]) # value doesn't matter + out = self.rbm.forward(v) + expected = -self.rbm.n_hidden * torch.log(torch.tensor(2.0)) + self.assertTrue(torch.allclose(out, expected)) + + def test_forward_ordering_bias(self): + """ + Free energy ordering test: + If visible_bias is very positive, visible=1 must yield + much lower free energy than visible=0. + """ + # Create a tiny RBM for easy testing + rbm = RBM(n_visible=1, n_hidden=1) + with torch.no_grad(): + rbm._weights[:] = 0 + rbm._hidden_biases[:] = 0 + rbm._visible_biases[:] = 1000.0 + + f1 = rbm.forward(torch.tensor([[1.0]])) + f0 = rbm.forward(torch.tensor([[0.0]])) + self.assertLess(f1, f0) + + def test_forward_small_numeric_case(self): + """Check free energy against manual calculation for 1 visible and 1 hidden unit.""" + rbm = RBM(n_visible=1, n_hidden=1) + with torch.no_grad(): + rbm._weights[:] = 2.0 + rbm._visible_biases[:] = 1.0 + rbm._hidden_biases[:] = -1.0 + + v = torch.tensor([[1.0]]) + # Independent manual calculation + expected = -1.0 - torch.nn.functional.softplus(torch.tensor(1.0)) + out = rbm.forward(v) + self.assertTrue(torch.allclose(out, expected)) + if __name__ == "__main__": unittest.main() From f8b9465c92e917dacc234ff640c999efd05c978b Mon Sep 17 00:00:00 2001 From: amansouribigvand Date: Fri, 12 Dec 2025 17:07:28 +0000 Subject: [PATCH 2/2] Decouple RBM, sampler, and optimizer + add release note --- .gitignore | 8 +- .../plugins/torch/models/boltzmann_machine.py | 194 ++---------------- dwave/plugins/torch/samplers/pcd_sampler.py | 58 ++++++ examples/rbm_image_generation.py | 179 ++++++++++------ .../notes/add-rbm-0b2134a1615ed5b3.yaml | 4 + tests/test_boltzmann_machine.py | 129 +----------- 6 files changed, 206 insertions(+), 366 deletions(-) create mode 100644 dwave/plugins/torch/samplers/pcd_sampler.py create mode 100644 releasenotes/notes/add-rbm-0b2134a1615ed5b3.yaml diff --git a/.gitignore b/.gitignore index bac9255..39c649a 100644 --- a/.gitignore +++ b/.gitignore @@ -66,10 +66,4 @@ venv.bak/ dmypy.json # aim -*.aim* - -# Datasets -data/* - -# Generated images -samples/* \ No newline at end of file +*.aim* \ No newline at end of file diff --git a/dwave/plugins/torch/models/boltzmann_machine.py b/dwave/plugins/torch/models/boltzmann_machine.py index e06f518..7f9bb44 100644 --- a/dwave/plugins/torch/models/boltzmann_machine.py +++ b/dwave/plugins/torch/models/boltzmann_machine.py @@ -666,8 +666,9 @@ def estimate_beta(self, spins: torch.Tensor) -> float: class RestrictedBoltzmannMachine(torch.nn.Module): """A Restricted Boltzmann Machine (RBM) model. - This class defines the parameterization and inference of a binary RBM. - Training is performed using Persistent Contrastive Divergence (PCD). + This class defines the parameterization of a binary RBM. + Training using Persistent Contrastive Divergence (PCD) must be + performed externally using separate sampler and optimizer classes. Args: n_visible (int): Number of visible units. @@ -699,14 +700,6 @@ def __init__( 0.5 * torch.ones(n_hidden) ) - # Stores the last visible states to initialize the Markov chain in Persistent Contrastive Divergence (PCD) - self.register_buffer("_previous_visible_values", None) - - # Initialize momenta tensors for momentum-based updates (all start at 0) - self.register_buffer("_weight_momenta", torch.zeros(n_visible, n_hidden)) - self.register_buffer("_visible_bias_momenta", torch.zeros(n_visible)) - self.register_buffer("_hidden_bias_momenta", torch.zeros(n_hidden)) - @property def n_visible(self) -> int: """Number of visible units.""" @@ -732,47 +725,26 @@ def hidden_biases(self) -> torch.Tensor: """Hidden biases of the RBM.""" return self._hidden_biases - @property - def previous_visible_values(self) -> torch.Tensor: - """Previous visible values used in Persistent Contrastive Divergence (PCD).""" - return self._previous_visible_values - - @property - def weight_momenta(self) -> torch.Tensor: - """Weight momenta of the RBM.""" - return self._weight_momenta - - @property - def visible_bias_momenta(self) -> torch.Tensor: - """Visible bias momenta of the RBM.""" - return self._visible_bias_momenta - - @property - def hidden_bias_momenta(self) -> torch.Tensor: - """Hidden bias momenta of the RBM.""" - return self._hidden_bias_momenta - - - def _sample_hidden(self, visible: torch.Tensor) -> torch.Tensor: + def sample_hidden(self, visible: torch.Tensor) -> torch.Tensor: """Sample from the distribution P(h|v). Args: visible (torch.Tensor): Tensor of shape (batch_size, n_visible) - representing the states of visible units. + representing the states of visible units. Returns: torch.Tensor: Binary tensor of shape (batch_size, n_hidden) representing - sampled hidden units. + sampled hidden units. """ - hidden_probs = torch.sigmoid(self._hidden_biases + visible @ self._weights) + hidden_probs = torch.sigmoid(self._hidden_biases + visible @ self._weights) return torch.bernoulli(hidden_probs) - def _sample_visible(self, hidden: torch.Tensor) -> torch.Tensor: + def sample_visible(self, hidden: torch.Tensor) -> torch.Tensor: """Sample from the distribution P(v|h). Args: hidden (torch.Tensor): Tensor of shape (batch_size, n_hidden) - representing the states of hidden units. + representing the states of hidden units. Returns: torch.Tensor: Binary tensor of shape (batch_size, n_visible) representing sampled visible units. @@ -780,155 +752,25 @@ def _sample_visible(self, hidden: torch.Tensor) -> torch.Tensor: visible_probs = torch.sigmoid(self._visible_biases + hidden @ self._weights.t()) return torch.bernoulli(visible_probs) - def generate_sample( - self, - batch_size: int, - gibbs_steps: int, - start_visible: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Generate a sample of visible and hidden units using gibbs sampling. - - Args: - batch_size (int): Number of samples to generate. - gibbs_steps (int): Number of Gibbs sampling steps to perform. - start_visible (torch.Tensor | None, optional): Initial visible states to - start the Gibbs chain (shape: [batch_size, n_visible]). If None, - a random Gaussian initialization is used. - - - Returns: - tuple[torch.Tensor, torch.Tensor]: A tuple of (visible, hidden) from the last Gibbs step: - - visible: (batch_size, n_visible) - - hidden: (batch_size, n_hidden) - """ - if start_visible is None: - visible_values = torch.randn( - batch_size, self.n_visible, device=self._weights.device - ) - else: - visible_values = start_visible - - hidden_values = None - - for _ in range(gibbs_steps): - hidden_values = self._sample_hidden(visible_values) - visible_values = self._sample_visible(hidden_values) - - return visible_values, hidden_values - - def _contrastive_divergence( - self, - batch: torch.Tensor, - epoch: int, - n_gibbs_steps: int, - learning_rate: float, - momentum_coefficient: float, - weight_decay: float, - n_epochs: int, - ) -> torch.Tensor: - """ - Perform one step of Contrastive Divergence (CD-k) with momentum and weight decay. - Uses Persistent Contrastive Divergence (PCD) by maintaining the last visible states - for Gibbs sampling across batches. - - Args: - batch (torch.Tensor): A batch of input data of shape (batch_size, n_visible). - epoch (int): Current training epoch. - n_gibbs_steps (int): Number of Gibbs sampling steps per epoch. - learning_rate (float): Base learning rate for parameter updates. - momentum_coefficient (float): Momentum coefficient for parameter updates. - weight_decay (float): weight decay (L2 regularization) coefficient for weights. - n_epochs (int): Number of training epochs. - - Returns: - torch.Tensor: The reconstruction error (L1 norm) for the batch. - - """ - - # Positive phase (data-driven) - hidden_probs = torch.sigmoid(self._hidden_biases + batch @ self._weights) - - weight_grads = torch.matmul(batch.t(), hidden_probs) - visible_bias_grads = batch - hidden_bias_grads = hidden_probs - - batch_size = batch.size(0) - - # Initialize previous visible states for Persistent CD - if self._previous_visible_values == None: - self._previous_visible_values = torch.randn_like( - batch, device=self._weights.device - ) - - # Negative phase (model-driven) - # Sample from the model using gibbs sampling - visible_values, hidden_values = self.generate_sample( - batch_size, n_gibbs_steps, self._previous_visible_values - ) - - visible_values = visible_values.detach() - hidden_values = hidden_values.detach() - # Store samples to initialize the next Markov chain with (PCD) - self._previous_visible_values = visible_values - - # Compute the gradients for negative phase - weight_grads -= torch.matmul(visible_values.t(), hidden_values) - - visible_bias_grads -= visible_values - hidden_bias_grads -= hidden_values - - # Average across the batch - weight_grads /= batch_size - visible_bias_grads /= batch_size - hidden_bias_grads /= batch_size - - # Compute decayed learning rate - decayed_learning_rate = learning_rate - (learning_rate / n_epochs * epoch) - - # Update momenta - self._weight_momenta = self._weight_momenta * momentum_coefficient + decayed_learning_rate * weight_grads - self._visible_bias_momenta = self._visible_bias_momenta * momentum_coefficient + decayed_learning_rate * torch.sum( - visible_bias_grads, dim=0 - ) - self._hidden_bias_momenta = self._hidden_bias_momenta * momentum_coefficient + decayed_learning_rate * torch.sum( - hidden_bias_grads, dim=0 - ) - - with torch.no_grad(): - # Update parameters - self._weights += self._weight_momenta - self._visible_biases += self._visible_bias_momenta - self._hidden_biases += self._hidden_bias_momenta - - # Apply weight decay - self._weights -= decayed_learning_rate * self._weights * weight_decay - - # Compute reconstruction error (L1 norm) - reconstruction = self._sample_visible(self._sample_hidden(batch)) - reconstruction = reconstruction.detach() - error = torch.sum(torch.abs(batch - reconstruction)) - - return error - - def forward(self, visible: torch.Tensor) -> torch.Tensor: - """ - Computes the RBM free energy of a batch of visible units averaged over the batch. + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Computes the RBM free energy of a batch of visible units averaged over the batch. - The free energy F(visible) for a visible vector visible is: + The free energy F(x) for a visible vector x is: - F(visible) = - visible · visible_biases - - sum_{j=1}^{n_hidden} log(1 + exp(hidden_biases[j] + (visible · weights)_j)) + .. math:: + F(x) = - x · visible_biases + - sum_{j=1}^{n_hidden} log(1 + exp(hidden_biases[j] + (x · weights)_j)) Args: - visible (torch.Tensor): Tensor of shape (batch_size, n_visible) representing the visible layer. + x (torch.Tensor): Tensor of shape (batch_size, n_visible) representing the visible layer. Returns: torch.Tensor: Scalar tensor representing the **average free energy** over the batch. """ - v_term = (visible * self._visible_biases).sum(dim=1) + v_term = (x * self._visible_biases).sum(dim=1) - hidden_pre_activation = visible @ self._weights + self._hidden_biases + hidden_pre_activation = x @ self._weights + self._hidden_biases h_term = torch.sum(torch.nn.functional.softplus(hidden_pre_activation), dim=1) diff --git a/dwave/plugins/torch/samplers/pcd_sampler.py b/dwave/plugins/torch/samplers/pcd_sampler.py new file mode 100644 index 0000000..364132b --- /dev/null +++ b/dwave/plugins/torch/samplers/pcd_sampler.py @@ -0,0 +1,58 @@ +import torch +from dwave.plugins.torch.models.boltzmann_machine import ( + RestrictedBoltzmannMachine as RBM, +) + +class PCDSampler: + """Persistent Contrastive Divergence (PCD) sampler for RBMs. + + This sampler maintains a persistent Markov chain of visible states + across minibatches and performs Gibbs sampling using the RBM’s + sampling functions. + + Args: + rbm (RBM): The RBM model from which the sampler draws samples. + """ + def __init__(self, rbm: RBM): + self.rbm = rbm + + # Stores the last visible states to initialize the Markov chain in Persistent Contrastive Divergence (PCD) + self.previous_visible_values = None + + def sample( + self, + batch_size: int, + gibbs_steps: int, + start_visible: torch.Tensor | None = None, + ): + """Generate a sample of visible and hidden units using gibbs sampling. + + Args: + batch_size (int): Number of samples to generate. + gibbs_steps (int): Number of Gibbs sampling steps to perform. + start_visible (torch.Tensor | None, optional): Initial visible states to + start the Gibbs chain (shape: [batch_size, n_visible]). If None, + a random Gaussian initialization is used. + + Returns: + tuple[torch.Tensor, torch.Tensor]: A tuple of (visible, hidden) from the last Gibbs step: + - visible: (batch_size, n_visible) + - hidden: (batch_size, n_hidden) + """ + if start_visible is None: + visible_values = torch.randn( + batch_size, self.rbm.n_visible, device=self.rbm.weights.device + ) + else: + visible_values = start_visible + + hidden_values = None + + for _ in range(gibbs_steps): + hidden_values = self.rbm.sample_hidden(visible_values) + visible_values = self.rbm.sample_visible(hidden_values) + + # Store samples to initialize the next Markov chain with (PCD) + self.previous_visible_values = visible_values.detach() + + return visible_values, hidden_values diff --git a/examples/rbm_image_generation.py b/examples/rbm_image_generation.py index bd5f14a..6e920d9 100644 --- a/examples/rbm_image_generation.py +++ b/examples/rbm_image_generation.py @@ -20,11 +20,12 @@ ) from torchvision import transforms, datasets import matplotlib.pyplot as plt +from torch.optim import SGD +from dwave.plugins.torch.samplers.pcd_sampler import PCDSampler def load_binarized_mnist(dataset_path: str = "data") -> datasets.MNIST: - """ - Load the MNIST dataset and binarize it (pixels >= 0.5 become 1, else 0). + """Load the MNIST dataset and binarize it (pixels >= 0.5 become 1, else 0). Args: dataset_path (str): Path to download/store the MNIST dataset. Defaults to "data". @@ -42,26 +43,90 @@ def load_binarized_mnist(dataset_path: str = "data") -> datasets.MNIST: return train_dataset +def contrastive_divergence( + rbm: RBM, + batch: torch.Tensor, + n_gibbs_steps: int, + sampler: PCDSampler, + optimizer: torch.optim.Optimizer, +) -> torch.Tensor: + """Perform one step of Contrastive Divergence (CD-k). + + Uses Persistent Contrastive Divergence (PCD) by maintaining the last visible states + for Gibbs sampling across batches. + Gradients are applied via the provided PyTorch optimizer. + + Args: + batch (torch.Tensor): A batch of input data of shape (batch_size, n_visible). + n_gibbs_steps (int): Number of Gibbs sampling steps per epoch. + sampler (PCDSampler): Sampler responsible for producing negative-phase samples. + optimizer (torch.optim.Optimizer): PyTorch optimizer. + + Returns: + torch.Tensor: The reconstruction error (L1 norm) for the batch. + """ + # Positive phase (data-driven) + hidden_probs = torch.sigmoid(batch @ rbm.weights + rbm.hidden_biases) + + weight_grads = torch.matmul(batch.t(), hidden_probs) + visible_bias_grads = batch.clone() + hidden_bias_grads = hidden_probs.clone() + + batch_size = batch.size(0) + + # Negative phase (model-driven) + # Sample from the model using gibbs sampling + visible_values, hidden_values = sampler.sample( + batch_size, + gibbs_steps=n_gibbs_steps, + start_visible=sampler.previous_visible_values, + ) + + visible_values = visible_values.detach() + hidden_values = hidden_values.detach() + + # Compute the gradients for negative phase + weight_grads -= torch.matmul(visible_values.t(), hidden_values) + visible_bias_grads -= visible_values + hidden_bias_grads -= hidden_values + + # Average across the batch + weight_grads /= batch_size + visible_bias_grads = torch.mean(visible_bias_grads, dim=0) + hidden_bias_grads = torch.mean(hidden_bias_grads, dim=0) + + # Apply gradients via optimizer + rbm.weights.grad = -weight_grads + rbm.visible_biases.grad = -visible_bias_grads + rbm.hidden_biases.grad = -hidden_bias_grads + + optimizer.step() + optimizer.zero_grad() + + # Compute reconstruction error (L1 norm) + reconstruction = rbm.sample_visible(rbm.sample_hidden(batch)) + reconstruction = reconstruction.detach() + reconstruction_error = torch.sum(torch.abs(batch - reconstruction)) + + return reconstruction_error + def train_loop( train_loader: DataLoader, rbm: RBM, n_epochs: int, n_gibbs_steps: int, - learning_rate: float, - momentum: float, - weight_decay: float, + sampler: PCDSampler, + optimizer: torch.optim.Optimizer, ) -> None: - """ - Train the RBM using contrastive divergence with momentum and weight decay. + """Train the RBM using contrastive divergence with a given PCDSampler and optimizer. Args: train_loader (DataLoader): PyTorch DataLoader for training data. rbm (RBM): Restricted Boltzmann Machine instance. n_epochs (int): Number of training epochs. n_gibbs_steps (int): Number of Gibbs sampling steps per CD update. - learning_rate (float): Base learning rate. - momentum (float): Momentum coefficient for parameter updates. - weight_decay (float): Weight decay (L2 regularization) coefficient. + sampler (PCDSampler): sampler responsible for producing negative-phase samples. + optimizer (torch.optim.Optimizer): PyTorch optimizer. """ device = rbm._weights.device for epoch in range(n_epochs): @@ -72,14 +137,8 @@ def train_loop( batch = batch.reshape(batch.size(0), rbm.n_visible).to(device) # Perform one step of contrastive divergence and accumulate error - error = rbm._contrastive_divergence( - batch, - epoch, - n_gibbs_steps, - learning_rate, - momentum, - weight_decay, - n_epochs, + error = contrastive_divergence( + rbm, batch, n_gibbs_steps, sampler, optimizer ) total_error += error num_examples += batch.size(0) @@ -90,18 +149,17 @@ def train_loop( def generate_and_save_images( - rbm: RBM, + sampler: PCDSampler, rows: int = 8, columns: int = 8, steps: int = 1000, output_dir: str = "samples", output_filename: str = "generated_images.png", -) -> None: - """ - Generate samples from a trained RBM and save them as a grid of images. +) -> PCDSampler: + """Generate samples from a trained RBM and save them as a grid of images. Args: - rbm (RBM): Trained RBM instance. + sampler (PCDSampler): sampler to generate samples from the trained RBM. rows (int): Number of rows in the output image grid. Defaults to 8. columns (int): Number of columns in the output image grid. Defaults to 8. steps (int): Number of Gibbs sampling steps for generation. Defaults to 1000. @@ -114,7 +172,10 @@ def generate_and_save_images( num_images = rows * columns # Generate batch of images - samples, _ = rbm.generate_sample(num_images, gibbs_steps=steps) + samples, _ = sampler.sample(num_images, gibbs_steps=steps) + + # for SpinRBM + #samples = ((samples + 1) / 2).view(num_images, 28, 28).detach().cpu().numpy() # convert -1/+1 → 0/1 samples = samples.view(num_images, 28, 28).detach().cpu().numpy() # Plot grid of images @@ -135,32 +196,30 @@ def generate_and_save_images( def train_rbm( - n_visible: int = 784, - n_hidden: int = 500, - n_gibbs_steps: int = 10, - learning_rate: float = 1e-3, - momentum: float = 0.5, - weight_decay: float = 1e-7, - n_epochs: int = 20, - batch_size: int = 32, + n_visible: int, + n_hidden: int, + n_gibbs_steps: int, + learning_rate: float, + momentum: float, + weight_decay: float, + n_epochs: int, + batch_size: int, dataset_path: str = "data", - output_dir: str = "samples", - output_filename: str = "generated_images.png", -) -> None: +) -> PCDSampler: """Train an RBM on MNIST and generate sample images. Args: - n_visible (int, optional): Number of visible units. Defaults to 784. - n_hidden (int, optional): Number of hidden units. Defaults to 500. - n_gibbs_steps (int, optional): Number of Gibbs sampling steps per CD update. Defaults to 10. - learning_rate (float, optional): Base learning rate for CD updates. Defaults to 1e-3. - momentum (float, optional): Momentum coefficient for CD updates. Defaults to 0.5. - weight_decay (float, optional): Weight decay (L2 regularization) coefficient. Defaults to 1e-7. - n_epochs (int, optional): Number of training epochs. Defaults to 20. - batch_size (int, optional): Batch size for training. Defaults to 32. + n_visible (int, optional): Number of visible units. + n_hidden (int, optional): Number of hidden units. + n_gibbs_steps (int, optional): Number of Gibbs sampling steps per CD update. + learning_rate (float, optional): Base learning rate for CD updates. + momentum (float, optional): Momentum coefficient for CD updates. + weight_decay (float, optional): Weight decay (L2 regularization) coefficient. + n_epochs (int, optional): Number of training epochs. + batch_size (int, optional): Batch size for training. dataset_path (str, optional): Path to download/store the MNIST dataset. Defaults to "data". - output_dir (str, optional): Directory to save the generated images. Defaults to "samples". - output_filename (str, optional): File name for saving the generated image grid. Defaults to "generated_images.png". + Returns: + PCDSampler: The sampler used for training the RBM. """ device = torch.device("cuda" if torch.cuda.is_available() else "cpu") @@ -178,31 +237,25 @@ def train_rbm( # Initialize RBM rbm = RBM(n_visible, n_hidden).to(device) - # Train RBM - train_loop( - train_loader, - rbm, - n_epochs=n_epochs, - n_gibbs_steps=n_gibbs_steps, - learning_rate=learning_rate, + # Initialize PCD Sampler + sampler = PCDSampler(rbm) + + optimizer = SGD( + [rbm.weights, rbm.visible_biases, rbm.hidden_biases], + lr=learning_rate, momentum=momentum, weight_decay=weight_decay, ) - # Generate and save samples - generate_and_save_images( - rbm, - rows=8, - columns=8, - steps=1000, - output_dir=output_dir, - output_filename=output_filename, - ) + # Train RBM + train_loop(train_loader, rbm, n_epochs, n_gibbs_steps, sampler, optimizer) + + return sampler if __name__ == "__main__": # Run an example of fitting a Restricted Boltzmann Machine to the MNIST dataset - train_rbm( + sampler = train_rbm( n_visible=784, n_hidden=500, n_gibbs_steps=10, @@ -212,3 +265,5 @@ def train_rbm( n_epochs=50, batch_size=64, ) + # Generate and save samples + generate_and_save_images(sampler) diff --git a/releasenotes/notes/add-rbm-0b2134a1615ed5b3.yaml b/releasenotes/notes/add-rbm-0b2134a1615ed5b3.yaml new file mode 100644 index 0000000..782dd56 --- /dev/null +++ b/releasenotes/notes/add-rbm-0b2134a1615ed5b3.yaml @@ -0,0 +1,4 @@ +--- +features: + - Add ``Restricted Boltzmann Machine`` class for training RBMs + using Persistant Contrastive Divergence algorithm. diff --git a/tests/test_boltzmann_machine.py b/tests/test_boltzmann_machine.py index 2e4b297..7e16b1e 100644 --- a/tests/test_boltzmann_machine.py +++ b/tests/test_boltzmann_machine.py @@ -429,13 +429,13 @@ def setUp(self): def test_sample_hidden_shape(self): visible = torch.randint(0, 2, (5, self.rbm.n_visible)).float() - hidden = self.rbm._sample_hidden(visible) + hidden = self.rbm.sample_hidden(visible) # Ensure shape is correct self.assertEqual(hidden.shape, (5, self.rbm.n_hidden)) def test_sample_hidden_binary(self): visible = torch.randint(0, 2, (5, self.rbm.n_visible)).float() - hidden = self.rbm._sample_hidden(visible) + hidden = self.rbm.sample_hidden(visible) # Ensure output is binary self.assertTrue(torch.all((hidden == 0) | (hidden == 1))) @@ -448,7 +448,7 @@ def test_sample_hidden_binary(self): ) def test_sample_hidden_saturation(self, name, bias_value, expected_value): """ - Test that _sample_hidden saturates correctly when the hidden biases + Test that sample_hidden saturates correctly when the hidden biases are set to very large positive or negative values. If hidden_bias[j] → +∞ @@ -468,21 +468,21 @@ def test_sample_hidden_saturation(self, name, bias_value, expected_value): visible = torch.zeros(5, self.rbm.n_visible) # Sample hidden units - hidden = self.rbm._sample_hidden(visible) + hidden = self.rbm.sample_hidden(visible) # Assert that all hidden units match the expected saturated value self.assertTrue(torch.all(hidden == expected_value)) def test_sample_visible_shape(self): hidden = torch.randint(0, 2, (5, self.rbm.n_hidden)).float() - visible = self.rbm._sample_visible(hidden) + visible = self.rbm.sample_visible(hidden) # Ensure shape is correct self.assertEqual(visible.shape, (5, self.rbm.n_visible)) def test_sample_visible_binary(self): hidden = torch.randint(0, 2, (5, self.rbm.n_hidden)).float() - visible = self.rbm._sample_visible(hidden) + visible = self.rbm.sample_visible(hidden) # Ensure output is binary self.assertTrue(torch.all((visible == 0) | (visible == 1)).item()) @@ -495,7 +495,7 @@ def test_sample_visible_binary(self): ) def test_sample_visible_saturation(self, name, bias_value, expected_value): """ - Test that _sample_visible saturates correctly when the visible biases + Test that sample_visible saturates correctly when the visible biases are set to very large positive or negative values. If visible_bias → +∞: sigmoid → 1 → bernoulli(1) → always 1 @@ -509,123 +509,10 @@ def test_sample_visible_saturation(self, name, bias_value, expected_value): # Hidden input doesn't matter when biases dominate hidden = torch.zeros(5, self.rbm.n_hidden) - visible = self.rbm._sample_visible(hidden) + visible = self.rbm.sample_visible(hidden) self.assertTrue(torch.all(visible == expected_value).item()) - def test_generate_sample_shape(self): - batch_size = 4 - visible, hidden = self.rbm.generate_sample(batch_size, gibbs_steps=2) - - # Ensure shapes are correct - self.assertEqual(visible.shape, (batch_size, self.rbm.n_visible)) - self.assertEqual(hidden.shape, (batch_size, self.rbm.n_hidden)) - - def test_generate_sample_binary(self): - batch_size = 4 - visible, hidden = self.rbm.generate_sample(batch_size, gibbs_steps=2) - - # Ensure outputs are binary - self.assertTrue(torch.all((visible == 0) | (visible == 1)).item()) - self.assertTrue(torch.all((hidden == 0) | (hidden == 1)).item()) - - def test_generate_sample_initial_hidden_start_visible(self): - """ - When providing a start_visible tensor and using gibbs_steps=1, the first - hidden sample should match _sample_hidden(start_visible). This ensures - the RBM correctly uses the provided initial visible state. - """ - batch_size = 4 - start = torch.zeros(batch_size, self.rbm.n_visible) - - torch.manual_seed(42) - _, hidden = self.rbm.generate_sample( - batch_size, gibbs_steps=1, start_visible=start - ) - - torch.manual_seed(42) - expected_hidden = self.rbm._sample_hidden(start) - - # First hidden sample should match _sample_hidden(start) - self.assertTrue(torch.equal(hidden, expected_hidden)) - - def test_generate_sample_more_gibbs_steps_changes_output(self): - batch_size = 4 - - torch.manual_seed(42) - v1, _ = self.rbm.generate_sample(batch_size, gibbs_steps=1) - - torch.manual_seed(42) - v4, _ = self.rbm.generate_sample(batch_size, gibbs_steps=4) - - # With same seed but longer chain, results should differ - self.assertFalse(torch.allclose(v1, v4)) - - def test_generate_sample_deterministic_small_rbm(self): - """ - Test generate_sample on a tiny deterministic RBM with known parameters. - This test uses a 3x3 RBM with manually set weights, hidden biases, and - visible biases to ensure that the outputs are fully predictable. - """ - rbm = RBM(n_visible=3, n_hidden=3) - - # Set deterministic weights and biases - rbm._weights.data = torch.tensor( - [[0.2, -0.1, 0.0], [0.1, 0.3, -0.2], [-0.2, 0.1, 0.2]] - ) - rbm._hidden_biases.data = torch.tensor([0.1, -0.2, 0.0]) - rbm._visible_biases.data = torch.tensor([0.5, 0.0, 0.3]) - - start_visible = torch.zeros(1, 3) - torch.manual_seed(42) - visible, hidden = rbm.generate_sample( - batch_size=1, gibbs_steps=1, start_visible=start_visible - ) - - # Expected values computed manually - expected_hidden = torch.tensor([[0.0, 0.0, 1.0]]) - expected_visible = torch.tensor([[0.0, 1.0, 1.0]]) - - # Check equality - self.assertTrue(torch.equal(hidden, expected_hidden)) - self.assertTrue(torch.equal(visible, expected_visible)) - - def test_contrastive_divergence_returns_nonnegative_tensor(self): - """CD returns a tensor of type torch.Tensor and L1 error >= 0""" - error = self.rbm._contrastive_divergence( - batch = self.batch, - epoch=0, - n_gibbs_steps=1, - learning_rate=0.01, - momentum_coefficient=0.5, - weight_decay=0.0, - n_epochs=1, - ) - - self.assertIsInstance(error, torch.Tensor) - self.assertGreaterEqual(error.item(), 0.0) - - @parameterized.expand( - [ - ("weights", "_weights"), - ("visible_biases", "_visible_biases"), - ("hidden_biases", "_hidden_biases"), - ] - ) - def test_cd_parameter_updates(self, name, attr): - """Test that the _contrastive_divergence function updates model parameters.""" - torch.manual_seed(42) - before = getattr(self.rbm, attr).clone() - - _ = self.rbm._contrastive_divergence(self.batch, **self.cd_kwargs) - - after = getattr(self.rbm, attr) - - # Check that parameters were updated - self.assertFalse( - torch.allclose(before, after), f"{attr} should update during CD" - ) - def test_forward_scalar_output(self): """Forward should return a scalar tensor.""" batch = torch.randn(5, self.rbm.n_visible)