From e08059377209a35ac536cfccc4f01a350ba78830 Mon Sep 17 00:00:00 2001 From: Shiro-Raven Date: Fri, 26 May 2023 23:22:35 +0200 Subject: [PATCH 1/3] added Rotosolve optimizer with tests --- lambeq/__init__.py | 6 +- lambeq/training/__init__.py | 2 + lambeq/training/rotosolve_optimizer.py | 171 +++++++++++++++++++++ tests/training/test_rotosolve_optimizer.py | 123 +++++++++++++++ 4 files changed, 300 insertions(+), 2 deletions(-) create mode 100644 lambeq/training/rotosolve_optimizer.py create mode 100644 tests/training/test_rotosolve_optimizer.py diff --git a/lambeq/__init__.py b/lambeq/__init__.py index 6dc64b8a..71bbb07c 100644 --- a/lambeq/__init__.py +++ b/lambeq/__init__.py @@ -85,6 +85,7 @@ 'Optimizer', 'SPSAOptimizer', + 'RotosolveOptimizer', 'Model', 'NumpyModel', @@ -126,8 +127,9 @@ stairs_reader, word_sequence_reader) from lambeq.tokeniser import Tokeniser, SpacyTokeniser from lambeq.training import (Checkpoint, Dataset, Optimizer, SPSAOptimizer, - Model, NumpyModel, PennyLaneModel, PytorchModel, - QuantumModel, TketModel, Trainer, PytorchTrainer, + RotosolveOptimizer, Model, NumpyModel, + PennyLaneModel, PytorchModel, QuantumModel, + TketModel, Trainer, PytorchTrainer, QuantumTrainer, BinaryCrossEntropyLoss, CrossEntropyLoss, LossFunction, MSELoss) from lambeq.version import (version as __version__, diff --git a/lambeq/training/__init__.py b/lambeq/training/__init__.py index 8b6da237..dc35e4ac 100644 --- a/lambeq/training/__init__.py +++ b/lambeq/training/__init__.py @@ -24,6 +24,7 @@ 'Optimizer', 'SPSAOptimizer', + 'RotosolveOptimizer', 'Trainer', 'PytorchTrainer', @@ -47,6 +48,7 @@ from lambeq.training.optimizer import Optimizer from lambeq.training.spsa_optimizer import SPSAOptimizer +from lambeq.training.rotosolve_optimizer import RotosolveOptimizer from lambeq.training.trainer import Trainer from lambeq.training.pytorch_trainer import PytorchTrainer diff --git a/lambeq/training/rotosolve_optimizer.py b/lambeq/training/rotosolve_optimizer.py new file mode 100644 index 00000000..22960eec --- /dev/null +++ b/lambeq/training/rotosolve_optimizer.py @@ -0,0 +1,171 @@ +# Copyright 2021-2023 Cambridge Quantum Computing Ltd. +# +# 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. + +""" +RotosolveOptimizer +============= +Module implementing the Rotosolve optimizer. + +""" +from __future__ import annotations + +from collections.abc import Callable, Iterable, Mapping +from typing import Any + +import numpy as np +from numpy.typing import ArrayLike + +from lambeq.training.optimizer import Optimizer +from lambeq.training.quantum_model import QuantumModel + + +class RotosolveOptimizer(Optimizer): + """An Optimizer using the Rotosolve algorithm. + + See https://quantum-journal.org/papers/q-2021-01-28-391/pdf/ for details. + + """ + + model : QuantumModel + + def __init__(self, model: QuantumModel, + hyperparams: dict[str, float], + loss_fn: Callable[[Any, Any], float], + bounds: ArrayLike | None = None) -> None: + """Initialise the Rotosolve optimizer. + + Parameters + ---------- + model : :py:class:`.QuantumModel` + A lambeq quantum model. + hyperparams : dict of str to float. + A dictionary containing the models hyperparameters. + loss_fn : Callable + A loss function of form `loss(prediction, labels)`. + bounds : ArrayLike, optional + The range of each of the model parameters. + + Raises + ------ + ValueError + If the length of `bounds` does not match the number + of the model parameters. + + """ + if bounds is None: + bounds = [[-np.pi, np.pi]]*len(model.weights) + + super().__init__(model, hyperparams, loss_fn, bounds) + + self.project: Callable[[np.ndarray], np.ndarray] + + bds = np.asarray(bounds) + if len(bds) != len(self.model.weights): + raise ValueError('Length of `bounds` must be the same as the ' + 'number of the model parameters') + self.project = lambda x: x.clip(bds[:, 0], bds[:, 1]) + + def backward( + self, + batch: tuple[Iterable[Any], np.ndarray]) -> float: + """Calculate the gradients of the loss function. + + The gradients are calculated with respect to the model + parameters. + + Parameters + ---------- + batch : tuple of Iterable and numpy.ndarray + Current batch. Contains an Iterable of diagrams in index 0, + and the targets in index 1. + + Returns + ------- + float + The calculated loss. + + """ + diagrams, targets = batch + + # The new model weights + self.gradient = np.copy(self.model.weights) + + old_model_weights = self.model.weights + + for i, _ in enumerate(self.gradient): + # Let phi be 0 + + # M_phi + self.gradient[i] = 0.0 + self.model.weights = self.gradient + m_phi = self.model(diagrams) + + # M_phi + pi/2 + self.gradient[i] = np.pi / 2 + self.model.weights = self.gradient + m_phi_plus = self.model(diagrams) + + # M_phi - pi/2 + self.gradient[i] = -np.pi / 2 + self.model.weights = self.gradient + m_phi_minus = self.model(diagrams) + + # Update weight + self.gradient[i] = -(np.pi / 2) - np.arctan2( + 2*m_phi - m_phi_plus - m_phi_minus, + m_phi_plus - m_phi_minus + ) + + # Calculate loss + self.model.weights = self.gradient + y1 = self.model(diagrams) + loss = self.loss_fn(y1, targets) + + self.model.weights = old_model_weights + + return loss + + def step(self) -> None: + """Perform optimisation step.""" + self.model.weights = self.gradient + self.model.weights = self.project(self.model.weights) + + self.update_hyper_params() + self.zero_grad() + + def update_hyper_params(self) -> None: + """Update the hyperparameters of the Rotosolve algorithm.""" + return + + def state_dict(self) -> dict[str, Any]: + """Return optimizer states as dictionary. + + Returns + ------- + dict + A dictionary containing the current state of the optimizer. + + """ + raise NotImplementedError + + def load_state_dict(self, state_dict: Mapping[str, Any]) -> None: + """Load state of the optimizer from the state dictionary. + + Parameters + ---------- + state_dict : dict + A dictionary containing a snapshot of the optimizer state. + + """ + raise NotImplementedError diff --git a/tests/training/test_rotosolve_optimizer.py b/tests/training/test_rotosolve_optimizer.py new file mode 100644 index 00000000..275fb86a --- /dev/null +++ b/tests/training/test_rotosolve_optimizer.py @@ -0,0 +1,123 @@ +import pytest + +import numpy as np + +from discopy import Cup, Word +from discopy.quantum.circuit import Id + +from lambeq import AtomicType, IQPAnsatz, RotosolveOptimizer + +N = AtomicType.NOUN +S = AtomicType.SENTENCE + +ansatz = IQPAnsatz({N: 1, S: 1}, n_layers=1, n_single_qubit_params=1) + +diagrams = [ + ansatz((Word("Alice", N) @ Word("runs", N >> S) >> Cup(N, N.r) @ Id(S))), + ansatz((Word("Alice", N) @ Word("walks", N >> S) >> Cup(N, N.r) @ Id(S))) +] + +from lambeq.training.model import Model + + +class ModelDummy(Model): + def __init__(self) -> None: + super().__init__() + self.initialise_weights() + def from_checkpoint(): + pass + def _make_lambda(self, diagram): + return diagram.lambdify(*self.symbols) + def initialise_weights(self): + self.weights = np.array([1.,2.,3.]) + def _clear_predictions(self): + pass + def _log_prediction(self, y): + pass + def get_diagram_output(self): + pass + def _make_checkpoint(self): + pass + def _load_checkpoint(self): + pass + def forward(self, x): + return self.weights.sum() + +loss = lambda yhat, y: np.abs(yhat-y).sum()**2 + +def test_init(): + model = ModelDummy.from_diagrams(diagrams) + model.initialise_weights() + optim = RotosolveOptimizer(model, + hyperparams={}, + loss_fn= loss, + bounds=[[-np.pi, np.pi]]*len(model.weights)) + + assert optim.project + +def test_backward(): + np.random.seed(3) + model = ModelDummy.from_diagrams(diagrams) + model.initialise_weights() + optim = RotosolveOptimizer(model, + hyperparams={}, + loss_fn= loss, + bounds=[[-np.pi, np.pi]]*len(model.weights)) + + optim.backward(([diagrams[0]], np.array([0]))) + + assert np.array_equal(optim.gradient.round(5), np.array([-1.5708] * len(model.weights))) + assert np.array_equal(model.weights, np.array([1.,2.,3.])) + +def test_step(): + np.random.seed(3) + model = ModelDummy.from_diagrams(diagrams) + model.initialise_weights() + optim = RotosolveOptimizer(model, + hyperparams={}, + loss_fn= loss, + bounds=[[-np.pi, np.pi]]*len(model.weights)) + optim.backward(([diagrams[0]], np.array([0]))) + optim.step() + + assert np.array_equal(model.weights.round(4), np.array([-1.5708] * len(model.weights))) + +def test_bound_error(): + model = ModelDummy() + model.initialise_weights() + with pytest.raises(ValueError): + _ = RotosolveOptimizer(model=model, + hyperparams={}, + loss_fn=loss, + bounds=[[0, 10]]*(len(model.weights)-1)) + +def test_none_bound_error(): + model = ModelDummy() + model.initialise_weights() + optim = RotosolveOptimizer(model=model, + hyperparams={}, + loss_fn=loss) + + assert optim.bounds == [[-np.pi, np.pi]] * len(model.weights) + +def test_load_state_dict(): + model = ModelDummy() + model.from_diagrams(diagrams) + model.initialise_weights() + optim = RotosolveOptimizer(model, + hyperparams={}, + loss_fn= loss) + + with pytest.raises(NotImplementedError): + optim.load_state_dict({}) + +def test_state_dict(): + model = ModelDummy() + model.from_diagrams(diagrams) + model.initialise_weights() + optim = RotosolveOptimizer(model, + hyperparams={}, + loss_fn= loss) + + with pytest.raises(NotImplementedError): + optim.state_dict() From 0940d7a8b290103148e93e3c9ed7b2ffae1587e9 Mon Sep 17 00:00:00 2001 From: Shiro-Raven Date: Sat, 3 Jun 2023 12:03:21 +0200 Subject: [PATCH 2/3] fix linting errors and apply PR style feedback to RotosolveOptimizer --- lambeq/__init__.py | 15 +++-- lambeq/training/__init__.py | 2 +- lambeq/training/rotosolve_optimizer.py | 23 +++---- tests/training/test_rotosolve_optimizer.py | 78 ++++++++++++++-------- 4 files changed, 68 insertions(+), 50 deletions(-) diff --git a/lambeq/__init__.py b/lambeq/__init__.py index 71bbb07c..c7444747 100644 --- a/lambeq/__init__.py +++ b/lambeq/__init__.py @@ -84,8 +84,8 @@ 'Dataset', 'Optimizer', - 'SPSAOptimizer', 'RotosolveOptimizer', + 'SPSAOptimizer', 'Model', 'NumpyModel', @@ -126,11 +126,12 @@ bag_of_words_reader, cups_reader, spiders_reader, stairs_reader, word_sequence_reader) from lambeq.tokeniser import Tokeniser, SpacyTokeniser -from lambeq.training import (Checkpoint, Dataset, Optimizer, SPSAOptimizer, - RotosolveOptimizer, Model, NumpyModel, - PennyLaneModel, PytorchModel, QuantumModel, - TketModel, Trainer, PytorchTrainer, - QuantumTrainer, BinaryCrossEntropyLoss, - CrossEntropyLoss, LossFunction, MSELoss) +from lambeq.training import (Checkpoint, Dataset, Optimizer, + RotosolveOptimizer, SPSAOptimizer, Model, + NumpyModel, PennyLaneModel, PytorchModel, + QuantumModel, TketModel, Trainer, + PytorchTrainer, QuantumTrainer, + BinaryCrossEntropyLoss, CrossEntropyLoss, + LossFunction, MSELoss) from lambeq.version import (version as __version__, version_tuple as __version_info__) diff --git a/lambeq/training/__init__.py b/lambeq/training/__init__.py index dc35e4ac..7cbba068 100644 --- a/lambeq/training/__init__.py +++ b/lambeq/training/__init__.py @@ -23,8 +23,8 @@ 'TketModel', 'Optimizer', - 'SPSAOptimizer', 'RotosolveOptimizer', + 'SPSAOptimizer', 'Trainer', 'PytorchTrainer', diff --git a/lambeq/training/rotosolve_optimizer.py b/lambeq/training/rotosolve_optimizer.py index 22960eec..6f8d287e 100644 --- a/lambeq/training/rotosolve_optimizer.py +++ b/lambeq/training/rotosolve_optimizer.py @@ -14,7 +14,7 @@ """ RotosolveOptimizer -============= +================== Module implementing the Rotosolve optimizer. """ @@ -33,8 +33,8 @@ class RotosolveOptimizer(Optimizer): """An Optimizer using the Rotosolve algorithm. - See https://quantum-journal.org/papers/q-2021-01-28-391/pdf/ for details. - + For detauls, check out: + https://quantum-journal.org/papers/q-2021-01-28-391/pdf/ """ model : QuantumModel @@ -73,7 +73,7 @@ def __init__(self, model: QuantumModel, bds = np.asarray(bounds) if len(bds) != len(self.model.weights): raise ValueError('Length of `bounds` must be the same as the ' - 'number of the model parameters') + 'number of the model parameters') self.project = lambda x: x.clip(bds[:, 0], bds[:, 1]) def backward( @@ -123,9 +123,9 @@ def backward( # Update weight self.gradient[i] = -(np.pi / 2) - np.arctan2( - 2*m_phi - m_phi_plus - m_phi_minus, + 2*m_phi - m_phi_plus - m_phi_minus, m_phi_plus - m_phi_minus - ) + ) # Calculate loss self.model.weights = self.gradient @@ -140,13 +140,8 @@ def step(self) -> None: """Perform optimisation step.""" self.model.weights = self.gradient self.model.weights = self.project(self.model.weights) - - self.update_hyper_params() - self.zero_grad() - def update_hyper_params(self) -> None: - """Update the hyperparameters of the Rotosolve algorithm.""" - return + self.zero_grad() def state_dict(self) -> dict[str, Any]: """Return optimizer states as dictionary. @@ -157,7 +152,7 @@ def state_dict(self) -> dict[str, Any]: A dictionary containing the current state of the optimizer. """ - raise NotImplementedError + raise {} def load_state_dict(self, state_dict: Mapping[str, Any]) -> None: """Load state of the optimizer from the state dictionary. @@ -168,4 +163,4 @@ def load_state_dict(self, state_dict: Mapping[str, Any]) -> None: A dictionary containing a snapshot of the optimizer state. """ - raise NotImplementedError + pass diff --git a/tests/training/test_rotosolve_optimizer.py b/tests/training/test_rotosolve_optimizer.py index 275fb86a..68e36c5d 100644 --- a/tests/training/test_rotosolve_optimizer.py +++ b/tests/training/test_rotosolve_optimizer.py @@ -24,100 +24,122 @@ class ModelDummy(Model): def __init__(self) -> None: super().__init__() self.initialise_weights() + def from_checkpoint(): pass + def _make_lambda(self, diagram): return diagram.lambdify(*self.symbols) + def initialise_weights(self): - self.weights = np.array([1.,2.,3.]) + self.weights = np.array([1., 2., 3.]) + def _clear_predictions(self): pass + def _log_prediction(self, y): pass + def get_diagram_output(self): pass + def _make_checkpoint(self): pass + def _load_checkpoint(self): pass + def forward(self, x): return self.weights.sum() + loss = lambda yhat, y: np.abs(yhat-y).sum()**2 + def test_init(): model = ModelDummy.from_diagrams(diagrams) model.initialise_weights() optim = RotosolveOptimizer(model, - hyperparams={}, - loss_fn= loss, - bounds=[[-np.pi, np.pi]]*len(model.weights)) + hyperparams={}, + loss_fn=loss, + bounds=[[-np.pi, np.pi]]*len(model.weights)) assert optim.project + def test_backward(): np.random.seed(3) model = ModelDummy.from_diagrams(diagrams) model.initialise_weights() optim = RotosolveOptimizer(model, - hyperparams={}, - loss_fn= loss, - bounds=[[-np.pi, np.pi]]*len(model.weights)) - + hyperparams={}, + loss_fn=loss, + bounds=[[-np.pi, np.pi]]*len(model.weights)) + optim.backward(([diagrams[0]], np.array([0]))) - assert np.array_equal(optim.gradient.round(5), np.array([-1.5708] * len(model.weights))) - assert np.array_equal(model.weights, np.array([1.,2.,3.])) + assert np.array_equal( + optim.gradient.round(5), + np.array([-1.5708] * len(model.weights)) + ) + assert np.array_equal(model.weights, np.array([1., 2., 3.])) + def test_step(): np.random.seed(3) model = ModelDummy.from_diagrams(diagrams) model.initialise_weights() optim = RotosolveOptimizer(model, - hyperparams={}, - loss_fn= loss, - bounds=[[-np.pi, np.pi]]*len(model.weights)) + hyperparams={}, + loss_fn=loss, + bounds=[[-np.pi, np.pi]]*len(model.weights)) optim.backward(([diagrams[0]], np.array([0]))) optim.step() - assert np.array_equal(model.weights.round(4), np.array([-1.5708] * len(model.weights))) + assert np.array_equal( + model.weights.round(4), + np.array([-1.5708] * len(model.weights)) + ) + def test_bound_error(): model = ModelDummy() model.initialise_weights() with pytest.raises(ValueError): _ = RotosolveOptimizer(model=model, - hyperparams={}, - loss_fn=loss, - bounds=[[0, 10]]*(len(model.weights)-1)) - + hyperparams={}, + loss_fn=loss, + bounds=[[0, 10]]*(len(model.weights)-1)) + + def test_none_bound_error(): model = ModelDummy() model.initialise_weights() optim = RotosolveOptimizer(model=model, - hyperparams={}, - loss_fn=loss) + hyperparams={}, + loss_fn=loss) assert optim.bounds == [[-np.pi, np.pi]] * len(model.weights) + def test_load_state_dict(): model = ModelDummy() model.from_diagrams(diagrams) model.initialise_weights() optim = RotosolveOptimizer(model, - hyperparams={}, - loss_fn= loss) - + hyperparams={}, + loss_fn=loss) + with pytest.raises(NotImplementedError): optim.load_state_dict({}) + def test_state_dict(): model = ModelDummy() model.from_diagrams(diagrams) model.initialise_weights() optim = RotosolveOptimizer(model, - hyperparams={}, - loss_fn= loss) - - with pytest.raises(NotImplementedError): - optim.state_dict() + hyperparams={}, + loss_fn=loss) + + assert optim.state_dict() == {} From c26bdbca755121ea42f18164a2fbb3d803626277 Mon Sep 17 00:00:00 2001 From: Shiro-Raven Date: Sat, 3 Jun 2023 12:05:42 +0200 Subject: [PATCH 3/3] adjust RotosolveOptimizer tests to new changes --- lambeq/training/rotosolve_optimizer.py | 2 +- tests/training/test_rotosolve_optimizer.py | 3 +-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/lambeq/training/rotosolve_optimizer.py b/lambeq/training/rotosolve_optimizer.py index 6f8d287e..9ff5f4fc 100644 --- a/lambeq/training/rotosolve_optimizer.py +++ b/lambeq/training/rotosolve_optimizer.py @@ -152,7 +152,7 @@ def state_dict(self) -> dict[str, Any]: A dictionary containing the current state of the optimizer. """ - raise {} + return {} def load_state_dict(self, state_dict: Mapping[str, Any]) -> None: """Load state of the optimizer from the state dictionary. diff --git a/tests/training/test_rotosolve_optimizer.py b/tests/training/test_rotosolve_optimizer.py index 68e36c5d..b451e5fa 100644 --- a/tests/training/test_rotosolve_optimizer.py +++ b/tests/training/test_rotosolve_optimizer.py @@ -130,8 +130,7 @@ def test_load_state_dict(): hyperparams={}, loss_fn=loss) - with pytest.raises(NotImplementedError): - optim.load_state_dict({}) + assert optim.load_state_dict({}) is None def test_state_dict():