Repository: IBM/aihwkit (language: python)

Issue:
Add new optimizer: Adam
## Description and motivation

As part of widening the scope of the library, it would be great to include more analog optimizers that can be used in different context and use cases. One of the likely first candidates would be one of the Adam variations due its popularity.

## Proposed solution

In a similar way to `AnalogSGD`, create a new subclass of the chosen `torch` optimizer, customizing its behavior when passing through the analog layers. 

## Alternatives and other information

* As this would be the second optimizer, some refactoring is also a likely candidate in order to generalize the common functionality for both optimizers (thank Corey for the suggestion!) - in particular, `regroup_param_groups()` seems like a good candidate for a generic function to be shared by all analog optimizers through a base class or similar means.
* Longer term, we should combine this approach with others that allow us to "adapt" other non-analog optimizers more easily - for example, making the tiles behave more like tensors or dividing the optimizers between analog and non-analog. These approaches are a bit more far-fetched


Discussion / hints:

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 91c9654a..8235e171 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -33,6 +33,8 @@ The format is based on [Keep a Changelog], and this project adheres to
   AnalogLinear. (\#240)
 * Return of weight gradients for ``InferenceTile`` (only),
   so that the gradient can be handled with any pytorch optimizer (\#241)
+* Added a generic analog optimizer `AnalogOptimizer` that allows extending
+  any existing optimizer with analog-specific features. (\#242)
 
 ### Changed
 
@@ -51,6 +53,8 @@ The format is based on [Keep a Changelog], and this project adheres to
 * Removed `weight` and `bias` of analog layers from the module
   parameters as these parameters are handled internally for analog
   tiles (\#241).
+* The `AnalogSGD` optimizer is now implemented based on the generic analog
+  optimizer, and its base module is `aihwkit.optim.analog_optimizer`. (\#242)
 
 ### Fixed
 
diff --git a/docs/source/api/aihwkit.optim.analog_inference.rst b/docs/source/api/aihwkit.optim.analog_inference.rst
new file mode 100644
index 00000000..dce27f9c
--- /dev/null
+++ b/docs/source/api/aihwkit.optim.analog_inference.rst
@@ -0,0 +1,8 @@
+aihwkit.optim.analog\_optimizer module
+======================================
+
+.. automodule:: aihwkit.optim.analog_optimizer
+   :members:
+   :undoc-members:
+   :show-inheritance:
+   :exclude-members: clone
diff --git a/docs/source/api/aihwkit.optim.analog_sgd.rst b/docs/source/api/aihwkit.optim.analog_sgd.rst
deleted file mode 100644
index f1caa9f0..00000000
--- a/docs/source/api/aihwkit.optim.analog_sgd.rst
+++ /dev/null
@@ -1,8 +0,0 @@
-aihwkit.optim.analog\_sgd module
-================================
-
-.. automodule:: aihwkit.optim.analog_sgd
-   :members:
-   :undoc-members:
-   :show-inheritance:
-   :exclude-members: clone
diff --git a/docs/source/api/aihwkit.optim.rst b/docs/source/api/aihwkit.optim.rst
index 4a9c66a6..0526e9d3 100644
--- a/docs/source/api/aihwkit.optim.rst
+++ b/docs/source/api/aihwkit.optim.rst
@@ -12,4 +12,4 @@ Submodules
 .. toctree::
    :maxdepth: 4
 
-   aihwkit.optim.analog_sgd
+   aihwkit.optim.analog_inference
diff --git a/docs/source/api_reference.rst b/docs/source/api_reference.rst
index 90cc1464..47cdb44a 100644
--- a/docs/source/api_reference.rst
+++ b/docs/source/api_reference.rst
@@ -25,7 +25,7 @@ API Reference
    aihwkit.nn.modules.conv
    aihwkit.nn.modules.linear
    aihwkit.optim
-   aihwkit.optim.analog_sgd
+   aihwkit.optim.analog_optimizer
    aihwkit.utils
    aihwkit.utils.visualization
 
diff --git a/src/aihwkit/nn/functions.py b/src/aihwkit/nn/functions.py
index d938366e..b4ad85d3 100644
--- a/src/aihwkit/nn/functions.py
+++ b/src/aihwkit/nn/functions.py
@@ -16,7 +16,7 @@
 
 from torch import Tensor, empty_like
 from torch.autograd import Function
-from aihwkit.simulator.tiles.base import AnalogContext
+from aihwkit.optim.context import AnalogContext
 
 
 class AnalogFunctionBase(Function):
diff --git a/src/aihwkit/optim/__init__.py b/src/aihwkit/optim/__init__.py
index b7334a8b..3df56b09 100644
--- a/src/aihwkit/optim/__init__.py
+++ b/src/aihwkit/optim/__init__.py
@@ -14,4 +14,4 @@
 
 # Convenience imports for easier access to the classes.
 
-from aihwkit.optim.analog_sgd import AnalogSGD
+from aihwkit.optim.analog_optimizer import AnalogOptimizer, AnalogSGD
diff --git a/src/aihwkit/optim/analog_sgd.py b/src/aihwkit/optim/analog_optimizer.py
similarity index 57%
rename from src/aihwkit/optim/analog_sgd.py
rename to src/aihwkit/optim/analog_optimizer.py
index ab407286..b77273a1 100644
--- a/src/aihwkit/optim/analog_sgd.py
+++ b/src/aihwkit/optim/analog_optimizer.py
@@ -10,27 +10,34 @@
 # copyright notice, and modified files need to carry a notice indicating
 # that they have been altered from the originals.
 
-"""Analog-aware stochastic gradient descent optimizer."""
+"""Analog-aware inference optimizer."""
 
-from typing import Callable, Optional
+from types import new_class
+from typing import Any, Callable, Dict, Optional, Type
 from warnings import warn
 
-from torch import clone, no_grad
+from torch.optim import Optimizer, SGD
 from torch.nn import Module
-from torch.optim import SGD
 
-from aihwkit.nn.modules.base import AnalogModuleBase
-from aihwkit.simulator.tiles.base import AnalogContext
+from aihwkit.optim.context import AnalogContext
 
 
-class AnalogSGD(SGD):
+class AnalogOptimizerMixin:
+    """Mixin for analog optimizers.
 
-    """Implements analog-aware stochastic gradient descent."""
+    This class contains the methods needed for enabling analog in an existing
+    ``Optimizer``. It is designed to be used as a mixin in conjunction with an
+    ``AnalogOptimizer`` or torch ``Optimizer``.
+    """
 
     def check_analog_module_devices(self, model: Module) -> None:
         """Checks and moves analog modules to the correct cuda device."""
-
         # TODO: remove this check and add a .to / .cuda to the AnalogContext
+
+        # Import dynamically, in order to avoid circular imports.
+        # pylint: disable=import-outside-toplevel
+        from aihwkit.nn.modules.base import AnalogModuleBase
+
         manual_cuda_move = False  # For showing only one warning.
         for (_, module) in model.named_modules():
             if isinstance(module, AnalogModuleBase):
@@ -61,7 +68,6 @@ def regroup_param_groups(self, model: Module) -> None:
         Args:
             model: model for the optimizer.
         """
-
         self.check_analog_module_devices(model)
 
         # Create the new param groups.
@@ -71,8 +77,8 @@ def regroup_param_groups(self, model: Module) -> None:
             rm_lst = []
             for param in group['params']:
                 if isinstance(param, AnalogContext):
-
-                    param.analog_tile.set_learning_rate(self.defaults['lr'])
+                    param.analog_tile.set_learning_rate(
+                        self.defaults['lr'])  # type: ignore[attr-defined]
                     analog_param_groups.append({
                         'params': [param],
                     })
@@ -86,11 +92,10 @@ def regroup_param_groups(self, model: Module) -> None:
         self.param_groups = [g for g in self.param_groups  # type: ignore[has-type]
                              if id(g) not in rm_group_lst]
 
-        # Add analog groups
+        # Add analog groups.
         for group in analog_param_groups:
-            self.add_param_group(group)
+            self.add_param_group(group)  # type: ignore[attr-defined]
 
-    @no_grad()
     def step(self, closure: Optional[Callable] = None) -> Optional[float]:
         """Performs an analog-aware single optimization step.
 
@@ -105,35 +110,32 @@ def step(self, closure: Optional[Callable] = None) -> Optional[float]:
         Returns:
             The loss, if ``closure`` has been passed as a parameter.
         """
-        # pylint: disable=too-many-branches,too-many-locals, protected-access
-        loss = None
-        if closure is not None:
-            loss = closure()
+        # pylint: disable=too-many-branches
+        # Update non-analog parameters using the given optimizer
+        ret = super().step(closure)  # type: ignore[misc]
 
+        # Update analog parameters
         for group in self.param_groups:
-            learning_rate = group['lr']
-            weight_decay = group['weight_decay']
-            momentum = group['momentum']
-            dampening = group['dampening']
-            nesterov = group['nesterov']
+            learning_rate = group.get('lr')
 
             # Use analog_tile object.
             for param in group['params']:
                 if isinstance(param, AnalogContext):
-                    # Handle internal analog update
+                    # Handle internal analog update.
                     analog_ctx = param
                     analog_tile = analog_ctx.analog_tile
 
                     if analog_ctx.use_torch_update:
-                        # in this case a separate weight parameter exists: do nothing
+                        # In this case a separate weight parameter exists: do nothing.
                         continue
 
-                    # Update learning rate
-                    analog_tile.set_learning_rate(learning_rate)
+                    # Update learning rate.
+                    if learning_rate is not None:
+                        analog_tile.set_learning_rate(learning_rate)
 
                     # Call `update` in the tile.
                     if not analog_ctx.has_gradient():
-                        # forward never used
+                        # Forward never used.
                         continue
 
                     if analog_ctx.use_indexed:
@@ -147,35 +149,14 @@ def step(self, closure: Optional[Callable] = None) -> Optional[float]:
 
                     analog_ctx.reset()
 
-                    continue
-
-                if param.grad is None:
-                    continue
-
-                d_p = param.grad
-                if weight_decay != 0:
-                    d_p = d_p.add(param, alpha=weight_decay)
-                if momentum != 0:
-                    param_state = self.state[param]
-                    if 'momentum_buffer' not in param_state:
-                        buf = param_state['momentum_buffer'] = clone(d_p).detach()
-                    else:
-                        buf = param_state['momentum_buffer']
-                        buf.mul_(momentum).add_(d_p, alpha=1 - dampening)
-                    if nesterov:
-                        d_p = d_p.add(buf, alpha=momentum)
-                    else:
-                        d_p = buf
-
-                param.add_(d_p, alpha=-group['lr'])
-
-            # Apply post-update step operations (diffuse, decay, etc).
-            # (only here because of unknown params order and shared weights)
+        # Apply post-update step operations (diffuse, decay, etc).
+        # (only here because of unknown params order and shared weights)
+        for group in self.param_groups:
             for param in group['params']:
                 if isinstance(param, AnalogContext):
                     param.analog_tile.post_update_step()
 
-        return loss
+        return ret
 
     def set_learning_rate(self, learning_rate: float = 0.1) -> None:
         """Update the learning rate to a new value.
@@ -192,3 +173,68 @@ def set_learning_rate(self, learning_rate: float = 0.1) -> None:
                 if isinstance(param, AnalogContext):
                     # Update learning rate on the tile
                     param.analog_tile.set_learning_rate(learning_rate)
+
+
+class AnalogOptimizer(AnalogOptimizerMixin, Optimizer):
+    """Generic optimizer that wraps an existing ``Optimizer`` for analog inference.
+
+    This class wraps an existing ``Optimizer``, customizing the optimization
+    step for triggering the analog update needed for analog tiles. All other
+    (digital) parameters are governed by the given torch optimizer. In case of
+    hardware-aware training (``InferenceTile``) the tile weight update is also
+    governed by the given optimizer, otherwise it is using the internal analog
+    update as defined in the ``rpu_config``.
+
+    The ``AnalogOptimizer`` constructor expects the wrapped optimizer class as
+    the first parameter, followed by any arguments required by the wrapped
+    optimizer.
+
+    Note:
+        The instances returned are of a *new* type that is a subclass of:
+
+        * the wrapped ``Optimizer`` (allowing access to all their methods and
+          attributes).
+        * this ``AnalogOptimizer``.
+
+    Example:
+        The following block illustrate how to create an optimizer that wraps
+        standard SGD:
+
+        >>> from torch.optim import SGD
+        >>> from torch.nn import Linear
+        >>> from aihwkit.simulator.configs.configs import InferenceRPUConfig
+        >>> from aihwkit.optim import AnalogOptimizer
+        >>> model = AnalogLinear(3, 4, rpu_config=InferenceRPUConfig)
+        >>> optimizer = AnalogOptimizer(SGD, model.parameters(), lr=0.02)
+    """
+
+    SUBCLASSES = {}  # type: Dict[str, Type]
+    """Registry of the created subclasses."""
+
+    def __new__(
+            cls,
+            optimizer_cls: Type,
+            *_: Any,
+            **__: Any
+    ) -> 'AnalogOptimizer':
+        subclass_name = '{}{}'.format(cls.__name__, optimizer_cls.__name__)
+
+        # Retrieve or create a new subclass, that inherits both from
+        # `AnalogOptimizer` and for the specific torch optimizer
+        # (`optimizer_cls`).
+        if subclass_name not in cls.SUBCLASSES:
+            cls.SUBCLASSES[subclass_name] = new_class(subclass_name, (cls, optimizer_cls), {})
+
+        return super().__new__(cls.SUBCLASSES[subclass_name])
+
+    def __init__(
+            self,
+            optimizer_cls: Type,  # pylint: disable=unused-argument
+            *args: Any,
+            **kwargs: Any
+    ):
+        super().__init__(*args, **kwargs)
+
+
+class AnalogSGD(AnalogOptimizerMixin, SGD):
+    """Implements analog-aware stochastic gradient descent."""
diff --git a/src/aihwkit/optim/context.py b/src/aihwkit/optim/context.py
new file mode 100644
index 00000000..c61590f0
--- /dev/null
+++ b/src/aihwkit/optim/context.py
@@ -0,0 +1,52 @@
+# -*- coding: utf-8 -*-
+
+# (C) Copyright 2020, 2021 IBM. All Rights Reserved.
+#
+# This code is licensed under the Apache License, Version 2.0. You may
+# obtain a copy of this license in the LICENSE.txt file in the root directory
+# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Any modifications or derivative works of this code must retain this
+# copyright notice, and modified files need to carry a notice indicating
+# that they have been altered from the originals.
+
+"""Parameter context for analog tiles."""
+
+from typing import Optional, Type, TYPE_CHECKING
+from torch import ones
+from torch.nn import Parameter
+
+if TYPE_CHECKING:
+    from aihwkit.simulator.tiles.base import BaseTile
+
+
+class AnalogContext(Parameter):
+    """Context for analog optimizer."""
+
+    # pylint: disable=signature-differs
+    def __new__(cls: Type['AnalogContext'], analog_tile: 'BaseTile') -> 'AnalogContext':
+        return Parameter.__new__(cls, data=ones((), device=analog_tile.device),
+                                 requires_grad=True)
+
+    def __init__(self, analog_tile: 'BaseTile'):
+        super().__init__()
+        self.analog_tile = analog_tile
+        self.use_torch_update = False
+        self.use_indexed = False
+        self.analog_input = []  # type: list
+        self.analog_grad_output = []  # type: list
+        self.shared_weights = None
+
+    def reset(self, analog_tile: Optional['BaseTile'] = None) -> None:
+        """Reset the gradient trace and optionally sets the tile pointer."""
+        if analog_tile is not None:
+            self.analog_tile = analog_tile
+        self.analog_input = []
+        self.analog_grad_output = []
+
+    def has_gradient(self) -> bool:
+        """Return whether a gradient trace was stored."""
+        return len(self.analog_input) > 0
+
+    def __repr__(self) -> str:
+        return 'AnalogContext of ' + self.analog_tile.get_brief_info()
diff --git a/src/aihwkit/simulator/tiles/base.py b/src/aihwkit/simulator/tiles/base.py
index c08ed754..0c496df3 100644
--- a/src/aihwkit/simulator/tiles/base.py
+++ b/src/aihwkit/simulator/tiles/base.py
@@ -13,54 +13,21 @@
 """High level analog tiles (base)."""
 
 from collections import OrderedDict
-from typing import Dict, Generic, List, Optional, Tuple, TypeVar, Union, Type
+from typing import Dict, Generic, List, Optional, Tuple, TypeVar, Union
 
 from numpy import concatenate, expand_dims
 from numpy import abs as numpy_abs
-from torch import Tensor, stack, ones
+from torch import Tensor, stack
 from torch import device as torch_device
-from torch.nn import Parameter
 from torch.autograd import no_grad
 
 from aihwkit.exceptions import TileError
 from aihwkit.simulator.rpu_base import tiles
-
+from aihwkit.optim.context import AnalogContext
 
 RPUConfigGeneric = TypeVar('RPUConfigGeneric')
 
 
-class AnalogContext(Parameter):
-    """Context for analog optimizer"""
-
-    # pylint: disable=signature-differs
-    def __new__(cls: Type['AnalogContext'], analog_tile: 'BaseTile') -> 'AnalogContext':
-        return Parameter.__new__(cls, data=ones((), device=analog_tile.device),
-                                 requires_grad=True)
-
-    def __init__(self, analog_tile: 'BaseTile'):
-        super().__init__()
-        self.analog_tile = analog_tile
-        self.use_torch_update = False
-        self.use_indexed = False
-        self.analog_input = []  # type: list
-        self.analog_grad_output = []  # type: list
-        self.shared_weights = None
-
-    def reset(self, analog_tile: Optional['BaseTile'] = None) -> None:
-        """ Resets the gradient trace and optionally sets the tile pointer. """
-        if analog_tile is not None:
-            self.analog_tile = analog_tile
-        self.analog_input = []
-        self.analog_grad_output = []
-
-    def has_gradient(self) -> bool:
-        """ Whether a gradient trace was stored """
-        return len(self.analog_input) > 0
-
-    def __repr__(self) -> str:
-        return 'AnalogContext of ' + self.analog_tile.get_brief_info()
-
-
 class BaseTile(Generic[RPUConfigGeneric]):
     """Base class for tiles.
 


Test patch:
diff --git a/tests/test_layers_lstm.py b/tests/test_layers_lstm.py
index 21d84819..bd8e3113 100644
--- a/tests/test_layers_lstm.py
+++ b/tests/test_layers_lstm.py
@@ -20,7 +20,7 @@
 from aihwkit.nn.modules.lstm import LSTMState
 from aihwkit.optim import AnalogSGD
 from aihwkit.simulator.configs.configs import InferenceRPUConfig
-from aihwkit.simulator.tiles.base import AnalogContext
+from aihwkit.optim.context import AnalogContext
 
 from .helpers.decorators import parametrize_over_layers
 from .helpers.layers import LSTM, LSTMCuda
diff --git a/tests/test_optimizers.py b/tests/test_optimizers.py
new file mode 100644
index 00000000..599984db
--- /dev/null
+++ b/tests/test_optimizers.py
@@ -0,0 +1,100 @@
+# -*- coding: utf-8 -*-
+
+# (C) Copyright 2020, 2021 IBM. All Rights Reserved.
+#
+# This code is licensed under the Apache License, Version 2.0. You may
+# obtain a copy of this license in the LICENSE.txt file in the root directory
+# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Any modifications or derivative works of this code must retain this
+# copyright notice, and modified files need to carry a notice indicating
+# that they have been altered from the originals.
+
+"""Tests for optimizers."""
+
+from torch import Tensor, manual_seed
+from torch.nn import Linear
+from torch.optim import SGD, AdamW
+from torch.nn.functional import mse_loss
+
+from aihwkit.optim import AnalogOptimizer
+
+from .helpers.testcases import AihwkitTestCase
+
+
+class InferenceOptimizerTest(AihwkitTestCase):
+    """Tests for the AnalogOptimizer."""
+
+    @staticmethod
+    def train_model(model, opt, loss_func, x_b, y_b):
+        """Train the model."""
+        epochs = 100
+        for _ in range(epochs):
+            opt.zero_grad()
+            pred = model(x_b)
+            loss = loss_func(pred, y_b)
+
+            loss.backward()
+            opt.step()
+
+        return loss
+
+    def test_instantiate(self):
+        """Test instantiating the optimizer."""
+        model = Linear(3, 4)
+        optimizer = AnalogOptimizer(SGD, model.parameters(), lr=0.123)
+
+        # Assert that a new subclass is created, with both the wrapped
+        # optimizer and the AnalogInferenceOptimizer as parents.
+        self.assertIsInstance(optimizer, AnalogOptimizer)
+        self.assertIsInstance(optimizer, SGD)
+        self.assertIsNot(type(optimizer), AnalogOptimizer)
+
+        # Assert over specific wrapped optimizer parameters.
+        self.assertEqual(optimizer.defaults['lr'], 0.123)
+
+    def test_train_digital_sgd(self):
+        """Test training digital layer with analog optimizer (SGD)."""
+        loss_func = mse_loss
+
+        x_b = Tensor([[0.1, 0.2, 0.3, 0.1], [0.2, 0.4, 0.3, 0.6]])
+        y_b = Tensor([[0.3], [0.6]])
+
+        manual_seed(4321)
+        model1 = Linear(4, 1)
+
+        manual_seed(4321)
+        model2 = Linear(4, 1)
+
+        learning_rate = 0.123
+        optimizer1 = AnalogOptimizer(SGD, model1.parameters(), lr=learning_rate)
+        optimizer2 = SGD(model2.parameters(), lr=learning_rate)
+
+        loss1 = self.train_model(model1, optimizer1, loss_func, x_b, y_b)
+        loss2 = self.train_model(model2, optimizer2, loss_func, x_b, y_b)
+
+        self.assertAlmostEqual(loss1, loss2)
+        self.assertTensorAlmostEqual(model1.weight, model2.weight)
+
+    def test_train_digital_adamw(self):
+        """Test training digital layer with analog optimizer (AdamW)."""
+        loss_func = mse_loss
+
+        x_b = Tensor([[0.1, 0.2, 0.3, 0.1], [0.2, 0.4, 0.3, 0.6]])
+        y_b = Tensor([[0.3], [0.6]])
+
+        manual_seed(4321)
+        model1 = Linear(4, 1)
+
+        manual_seed(4321)
+        model2 = Linear(4, 1)
+
+        learning_rate = 0.123
+        optimizer1 = AnalogOptimizer(AdamW, model1.parameters(), lr=learning_rate)
+        optimizer2 = AdamW(model2.parameters(), lr=learning_rate)
+
+        loss1 = self.train_model(model1, optimizer1, loss_func, x_b, y_b)
+        loss2 = self.train_model(model2, optimizer2, loss_func, x_b, y_b)
+
+        self.assertAlmostEqual(loss1, loss2)
+        self.assertTensorAlmostEqual(model1.weight, model2.weight)