From 11331f75694f925e217505392f74a0cf208b1b09 Mon Sep 17 00:00:00 2001 From: epapoutsellis Date: Tue, 23 Aug 2022 21:07:43 +0100 Subject: [PATCH] fix SGD #1345, update SAG, SAGA --- .../optimisation/functions/SAGAFunction.py | 19 +++++-------------- .../cil/optimisation/functions/SAGFunction.py | 18 ++++-------------- .../cil/optimisation/functions/SGDFunction.py | 18 ++++-------------- Wrappers/Python/test/test_SAG_function.py | 4 ++-- 4 files changed, 15 insertions(+), 44 deletions(-) diff --git a/Wrappers/Python/cil/optimisation/functions/SAGAFunction.py b/Wrappers/Python/cil/optimisation/functions/SAGAFunction.py index faf4ba9481..66f2ffc199 100644 --- a/Wrappers/Python/cil/optimisation/functions/SAGAFunction.py +++ b/Wrappers/Python/cil/optimisation/functions/SAGAFunction.py @@ -54,7 +54,7 @@ def __init__(self, functions, sampling = "random", precond=None, replacement = F self.allocate_memory = False super(SAGFunction, self).__init__(functions, sampling = sampling, replacement=replacement) - def gradient(self, x, out=None): + def gradient(self, x, out): """ Returns a variance-reduced approximate gradient, defined below. For f = 1/num_subsets \sum_{i=1}^num_subsets F_{i}, the output is computed as follows: @@ -98,26 +98,17 @@ def gradient(self, x, out=None): self.tmp1.sapyb(1., self.subset_gradients[self.subset_num], -1., out=self.tmp2) # Compute the output : tmp2 + full_gradient - if out is None: - ret = 0.0 * self.tmp2 - self.tmp2.sapyb(1., self.full_gradient, 1., out=ret) - else: - self.tmp2.sapyb(1., self.full_gradient, 1., out=out) + self.tmp2.sapyb(1., self.full_gradient, 1., out=out) # Apply preconditioning - if self.precond is not None: - if out is None: - ret.multiply(self.precond,out=ret) - else: - out.multiply(self.precond,out=out) + if self.precond is not None: + out.multiply(self.precond(self.subset_num, x), out=out) # Update subset gradients in memory: store the computed gradient F_{subset_num} (x) in self.subset_gradients[self.subset_num] self.subset_gradients[self.subset_num].fill(self.tmp1) # Update the full gradient estimator: add 1/num_subsets * (gradient F_{subset_num} (x) - subset_gradient_in_memory_{subset_num}) to the current full_gradient self.full_gradient.sapyb(1., self.tmp2, 1./self.num_subsets, out=self.full_gradient) - - if out is None: - return ret + \ No newline at end of file diff --git a/Wrappers/Python/cil/optimisation/functions/SAGFunction.py b/Wrappers/Python/cil/optimisation/functions/SAGFunction.py index 8dd036da01..80fa866ec5 100644 --- a/Wrappers/Python/cil/optimisation/functions/SAGFunction.py +++ b/Wrappers/Python/cil/optimisation/functions/SAGFunction.py @@ -54,7 +54,7 @@ def __init__(self, functions, sampling = "random", precond=None, replacement = F self.allocate_memory = False super(SAGFunction, self).__init__(functions, sampling = sampling, replacement=replacement) - def gradient(self, x, out=None): + def gradient(self, x, out): """ Returns a variance-reduced approximate gradient. @@ -78,18 +78,11 @@ def gradient(self, x, out=None): self.tmp1.sapyb(1., self.subset_gradients[self.subset_num], -1., out=self.tmp2) # Compute the output : 1/num_subsets * tmp2 + full_gradient - if out is None: - ret = 0.0 * self.tmp2 - self.tmp2.sapyb(1./self.num_subsets, self.full_gradient, 1., out=ret) - else: - self.tmp2.sapyb(1./self.num_subsets, self.full_gradient, 1., out=out) + self.tmp2.sapyb(1./self.num_subsets, self.full_gradient, 1., out=out) # Apply preconditioning - if self.precond is not None: - if out is None: - ret.multiply(self.precond,out=ret) - else: - out.multiply(self.precond,out=out) + if self.precond is not None: + out.multiply(self.precond(self.subset_num, x), out=out) # Update subset gradients in memory: store the computed gradient F_{subset_num} (x) in self.subset_gradients[self.subset_num] self.subset_gradients[self.subset_num].fill(self.tmp1) @@ -97,9 +90,6 @@ def gradient(self, x, out=None): # Update the full gradient estimator: add 1/num_subsets * (gradient F_{subset_num} (x) - subset_gradient_in_memory_{subset_num}) to the current full_gradient self.full_gradient.sapyb(1., self.tmp2, 1./self.num_subsets, out=self.full_gradient) - if out is None: - return ret - def initialise_memory(self, x): r"""Initialize subset gradients :math:`v_{i}` and full gradient that are stored in memory. diff --git a/Wrappers/Python/cil/optimisation/functions/SGDFunction.py b/Wrappers/Python/cil/optimisation/functions/SGDFunction.py index ac431a8eae..9828130274 100644 --- a/Wrappers/Python/cil/optimisation/functions/SGDFunction.py +++ b/Wrappers/Python/cil/optimisation/functions/SGDFunction.py @@ -51,7 +51,7 @@ def __init__(self, functions, sampling = "random", replacement = False, precond= super(SGDFunction, self).__init__(functions, sampling = sampling, replacement = replacement) self.precond = precond - def gradient(self, x, out=None): + def gradient(self, x, out): """ Returns the gradient of the selected function at :code:`x`. The function is selected using the :meth:`~SubsetSumFunction.next_subset` """ @@ -60,19 +60,9 @@ def gradient(self, x, out=None): self.next_subset() # Compute new gradient for current subset - if out is None: - ret = 0.0 * x - self.functions[self.subset_num].gradient(x, out=ret) - else: - self.functions[self.subset_num].gradient(x, out=out) + self.functions[self.subset_num].gradient(x, out=out) # Apply preconditioning - if self.precond is not None: - if out is None: - ret.multiply(self.precond,out=ret) - else: - out.multiply(self.precond,out=out) - - if out is None: - return ret + if self.precond is not None: + out.multiply(self.precond(self.subset_num, x), out=out) \ No newline at end of file diff --git a/Wrappers/Python/test/test_SAG_function.py b/Wrappers/Python/test/test_SAG_function.py index dc0bc59512..c95189f50e 100644 --- a/Wrappers/Python/test/test_SAG_function.py +++ b/Wrappers/Python/test/test_SAG_function.py @@ -45,7 +45,7 @@ def setUp(self): self.F = LeastSquares(self.Aop, b=self.bop, c = 0.5/self.n_subsets) self.ig = self.Aop.domain - self.precond = self.ig.allocate(1.0) + self.precond = lambda i, x: 3./self.ig.allocate(2.5) self.F_SAG = SAGFunction(self.fi_cil, replacement = True, precond = self.precond) self.initial = self.ig.allocate() @@ -69,7 +69,7 @@ def test_gradient(self): tmp_sag.functions[tmp_sag.subset_num].gradient(x, out=tmp_sag.tmp1) tmp_sag.tmp1.sapyb(1., tmp_sag.subset_gradients[tmp_sag.subset_num], -1., out=tmp_sag.tmp2) tmp_sag.tmp2.sapyb(1./tmp_sag.num_subsets, tmp_sag.full_gradient, 1., out=out2) - out2.multiply(self.precond,out=out2) + out2 *= self.precond(tmp_sag.tmp2.subset_num, 3./self.ig.allocate(2.5)) # update subset_gradient in the subset_num # update full gradient