Adding learnable output scaling alpha (#353)

Author: Fabio-83

CHANGELOG.md

@@ -12,7 +12,7 @@ The format is based on [Keep a Changelog], and this project adheres to
 * `Fixed` for any bug fixes.
 * `Security` in case of vulnerabilities.
 
-## [0.5.0] - 2022/01/19
+## [0.5.0] - 2022/01/27
 
 ### Added
 
@@ -36,6 +36,8 @@ The format is based on [Keep a Changelog], and this project adheres to
   output sizes can be configured for the ``*Mapped`` layers. (\#331)
 * Notebooks directory with several notebook examples (#333, \#334)
 * Analog information summary function. (\#316)
+* The `alpha` weight scaling factor can now be defined as learnable parameter by switching
+  `learn_out_scaling_alpha` in the `rpu_config.mapping` parameters. (\#353)
 
 ### Fixed
 
@@ -62,6 +64,7 @@ The format is based on [Keep a Changelog], and this project adheres to
 * Digital bias is now accessable through ``MappingParameter``. (\#331)
 * The aihwkit documentation. New content around analog ai concepts, training presets, analog ai 
   optimizers, new references, and examples. (\#348)
+* The `weight_scaling_omega` can now be defined in the `rpu_config.mapping`. (\#353)
 
 ### Deprecated
 

examples/11_vgg8_training.py

@@ -35,6 +35,7 @@
 from aihwkit.optim import AnalogSGD
 from aihwkit.simulator.presets.configs import GokmenVlasovPreset
 from aihwkit.simulator.rpu_base import cuda
+from aihwkit.simulator.configs.utils import MappingParameter
 
 # Check device
 USE_CUDA = 0
@@ -59,7 +60,8 @@
 
 # Select the device model to use in the training. In this case we are using one of the preset,
 # but it can be changed to a number of preset to explore possible different analog devices
-RPU_CONFIG = GokmenVlasovPreset()
+mapping = MappingParameter(weight_scaling_omega=WEIGHT_SCALING_OMEGA)
+RPU_CONFIG = GokmenVlasovPreset(mapping=mapping)
 
 
 def load_images():
@@ -93,33 +95,33 @@ def create_analog_network():
         nn.ReLU(),
         AnalogConv2d(in_channels=channel[0], out_channels=channel[0], kernel_size=3, stride=1,
                      padding=1,
-                     rpu_config=RPU_CONFIG, weight_scaling_omega=WEIGHT_SCALING_OMEGA),
+                     rpu_config=RPU_CONFIG),
         nn.BatchNorm2d(channel[0]),
         nn.ReLU(),
         nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1),
         AnalogConv2d(in_channels=channel[0], out_channels=channel[1], kernel_size=3, stride=1,
                      padding=1,
-                     rpu_config=RPU_CONFIG, weight_scaling_omega=WEIGHT_SCALING_OMEGA),
+                     rpu_config=RPU_CONFIG),
         nn.ReLU(),
         AnalogConv2d(in_channels=channel[1], out_channels=channel[1], kernel_size=3, stride=1,
                      padding=1,
-                     rpu_config=RPU_CONFIG, weight_scaling_omega=WEIGHT_SCALING_OMEGA),
+                     rpu_config=RPU_CONFIG),
         nn.BatchNorm2d(channel[1]),
         nn.ReLU(),
         nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1),
         AnalogConv2d(in_channels=channel[1], out_channels=channel[2], kernel_size=3, stride=1,
                      padding=1,
-                     rpu_config=RPU_CONFIG, weight_scaling_omega=WEIGHT_SCALING_OMEGA),
+                     rpu_config=RPU_CONFIG),
         nn.ReLU(),
         AnalogConv2d(in_channels=channel[2], out_channels=channel[2], kernel_size=3, stride=1,
                      padding=1,
-                     rpu_config=RPU_CONFIG, weight_scaling_omega=WEIGHT_SCALING_OMEGA),
+                     rpu_config=RPU_CONFIG),
         nn.BatchNorm2d(channel[2]),
         nn.ReLU(),
         nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1),
         nn.Flatten(),
         AnalogLinear(in_features=16 * channel[2], out_features=fc_size,
-                     rpu_config=RPU_CONFIG, weight_scaling_omega=WEIGHT_SCALING_OMEGA),
+                     rpu_config=RPU_CONFIG),
         nn.ReLU(),
         nn.Linear(in_features=fc_size, out_features=N_CLASSES),
         nn.LogSoftmax(dim=1)

examples/16_mnist_gan.py

@@ -52,8 +52,9 @@
 # optimizer used (e.g. mixed precision or full analog update)
 
 # As an example we use a mixed precision preset using an ECRAM device model
-RPU_CONFIG = MixedPrecisionEcRamMOPreset()
-WS_OMEGA = 0.8
+from aihwkit.simulator.configs.utils import MappingParameter
+mapping = MappingParameter(weight_scaling_omega=0.8)
+RPU_CONFIG = MixedPrecisionEcRamMOPreset(mapping=mapping)
 
 # Set your parameters
 SEED = 1
@@ -129,8 +130,7 @@ def get_generator_block(input_dim, output_dim):
             input_dim,
             output_dim,
             bias=True,
-            rpu_config=RPU_CONFIG,
-            weight_scaling_omega=WS_OMEGA,
+            rpu_config=RPU_CONFIG
         ),
         nn.BatchNorm1d(output_dim),
         nn.ReLU(inplace=True),
@@ -159,8 +159,7 @@ def __init__(self, z_dim=10, im_dim=784, hidden_dim=128):
                 hidden_dim * 8,
                 im_dim,
                 bias=True,
-                rpu_config=RPU_CONFIG,
-                weight_scaling_omega=WS_OMEGA,
+                rpu_config=RPU_CONFIG
             ),
             nn.Sigmoid(),
         )

examples/17_resnet34_imagenet_conversion_to_analog.py

@@ -35,13 +35,14 @@
 # Define device and chip configuration used in the RPU tile
 mapping = MappingParameter(max_input_size=512,  # analog tile size
                            max_output_size=512,
-                           digital_bias=True)  # whether to use analog or digital bias
+                           digital_bias=True,
+                           weight_scaling_omega=0.6)  # whether to use analog or digital bias
 # Choose any preset or RPU configuration here
 rpu_config = TikiTakaReRamSBPreset(mapping=mapping)
 
 # Convert the model to its analog version.
 # this will replace ``Linear`` layers with ``AnalogLinearMapped``
-model = convert_to_analog_mapped(model, rpu_config, weight_scaling_omega=0.6)
+model = convert_to_analog_mapped(model, rpu_config)
 
 # Note: One can also use ``convert_to_analog`` instead to convert
 # ``Linear`` to ``AnalogLinear`` (without mapping to multiple tiles)

examples/18_cifar10_on_resnet.py

@@ -34,6 +34,7 @@
 from aihwkit.nn.conversion import convert_to_analog
 from aihwkit.simulator.presets import TikiTakaEcRamPreset
 from aihwkit.simulator.rpu_base import cuda
+from aihwkit.simulator.configs.utils import MappingParameter
 
 # Device to use
 USE_CUDA = 0
@@ -56,7 +57,8 @@
 N_CLASSES = 10
 
 # Device used in the RPU tile
-RPU_CONFIG = TikiTakaEcRamPreset()
+mapping = MappingParameter(weight_scaling_omega=0.6)
+RPU_CONFIG = TikiTakaEcRamPreset(mapping=mapping)
 
 
 class ResidualBlock(nn.Module):
@@ -295,7 +297,7 @@ def main():
     model = create_model()
 
     # Convert the model to its analog version
-    model = convert_to_analog(model, RPU_CONFIG, weight_scaling_omega=0.6)
+    model = convert_to_analog(model, RPU_CONFIG)
     # Load saved weights if previously saved
     # model.load_state_dict(load(WEIGHT_PATH))
 

src/aihwkit/nn/conversion.py

@@ -45,7 +45,6 @@ def convert_to_analog(
         module: Module,
         rpu_config: RPUConfigGeneric,
         realistic_read_write: bool = False,
-        weight_scaling_omega: float = 0.0,
         conversion_map: Optional[Dict] = None
 ) -> Module:
     """Convert a given digital model to analog counter parts.
@@ -63,9 +62,6 @@ def convert_to_analog(
             Applied to all converted tiles.
         realistic_read_write: Whether to use closed-loop programming
             when setting the weights. Applied to all converted tiles.
-        weight_scaling_omega: If non-zero, the analog weights will be
-            scaled by ``weight_scaling_omega`` divided by the absolute
-            maximum value of the original weight matrix.
 
             Note:
                 Make sure that the weight max and min settings of the
@@ -92,7 +88,7 @@ def convert_to_analog(
     # Convert parent.
     if module.__class__ in conversion_map:
         module = conversion_map[module.__class__].from_digital(  # type: ignore
-            module, rpu_config, realistic_read_write, weight_scaling_omega)
+            module, rpu_config, realistic_read_write)
 
     # Convert children.
     convert_dic = {}
@@ -101,11 +97,11 @@ def convert_to_analog(
         n_grand_children = len(list(mod.named_children()))
         if n_grand_children > 0:
             new_mod = convert_to_analog(mod, rpu_config, realistic_read_write,
-                                        weight_scaling_omega, conversion_map)
+                                        conversion_map)
 
         elif mod.__class__ in conversion_map:
             new_mod = conversion_map[mod.__class__].from_digital(   # type: ignore
-                mod, rpu_config, realistic_read_write, weight_scaling_omega)
+                mod, rpu_config, realistic_read_write)
         else:
             continue
 
@@ -125,7 +121,6 @@ def convert_to_analog_mapped(
         module: Module,
         rpu_config: RPUConfigGeneric,
         realistic_read_write: bool = False,
-        weight_scaling_omega: float = 0.0,
 ) -> Module:
     """Convert a given digital model to its analog counterpart with tile
     mapping support.
@@ -142,9 +137,6 @@ def convert_to_analog_mapped(
         rpu_config: RPU config to apply to all converted tiles.
         realistic_read_write: Whether to use closed-loop programming
             when setting the weights. Applied to all converted tiles.
-        weight_scaling_omega: If non-zero, the analog weights will be
-            scaled by ``weight_scaling_omega`` divided by the absolute
-            maximum value of the original weight matrix.
 
             Note:
                 Make sure that the weight max and min settings of the
@@ -159,6 +151,5 @@ def convert_to_analog_mapped(
         module,
         rpu_config,
         realistic_read_write,
-        weight_scaling_omega,
         _DEFAULT_MAPPED_CONVERSION_MAP
     )

src/aihwkit/nn/modules/base.py

@@ -11,6 +11,7 @@
 # that they have been altered from the originals.
 
 """Base class for analog Modules."""
+import warnings
 
 from typing import (
     Any, Dict, List, Optional, Tuple, NamedTuple, Union,
@@ -72,14 +73,15 @@ class AnalogModuleBase(Module):
     ANALOG_CTX_PREFIX: str = 'analog_ctx_'
     ANALOG_SHARED_WEIGHT_PREFIX: str = 'analog_shared_weights_'
     ANALOG_STATE_PREFIX: str = 'analog_tile_state_'
+    ANALOG_OUT_SCALING_ALPHA_PREFIX: str = 'analog_out_scaling_alpha_'
 
     def __init__(
             self,
             in_features: int,
             out_features: int,
             bias: bool,
             realistic_read_write: bool = False,
-            weight_scaling_omega: float = 0.0,
+            weight_scaling_omega: Optional[float] = None,
             mapping: Optional[MappingParameter] = None,
     ) -> None:
         # pylint: disable=super-init-not-called
@@ -93,9 +95,21 @@ def __init__(
         self.use_bias = bias
         self.digital_bias = bias and mapping.digital_bias
         self.analog_bias = bias and not mapping.digital_bias
+        self.weight_scaling_omega = mapping.weight_scaling_omega if weight_scaling_omega is None \
+            else weight_scaling_omega
+        if weight_scaling_omega is not None:
+            warnings.warn(DeprecationWarning('\nSetting the weight_scaling_omega through the '
+                                             'layers input parameters will be deprecated in the '
+                                             'future. Please set it through the MappingParameter '
+                                             'of the rpu_config.\n'))
+
+        self.weight_scaling_omega_columnwise = mapping.weight_scaling_omega_columnwise
+        self.learn_out_scaling_alpha = mapping.learn_out_scaling_alpha
+
+        if self.learn_out_scaling_alpha and self.weight_scaling_omega == 0:
+            raise ValueError('out_scaling_alpha can only be learned if weight_scaling_omega > 0')
 
         self.realistic_read_write = realistic_read_write
-        self.weight_scaling_omega = weight_scaling_omega
         self.in_features = in_features
         self.out_features = out_features
 
@@ -129,6 +143,15 @@ def register_analog_tile(self, tile: 'BaseTile', name: Optional[str] = None) ->
             if par_name not in self._registered_helper_parameter:
                 self._registered_helper_parameter.append(par_name)
 
+        if self.learn_out_scaling_alpha:
+            if not isinstance(tile.out_scaling_alpha, Parameter):
+                tile.out_scaling_alpha = Parameter(tile.out_scaling_alpha)
+            par_name = self.ANALOG_OUT_SCALING_ALPHA_PREFIX + str(self._analog_tile_counter)
+            self.register_parameter(par_name, tile.out_scaling_alpha)
+
+            if par_name not in self._registered_helper_parameter:
+                self._registered_helper_parameter.append(par_name)
+
         self._analog_tile_counter += 1
 
     def unregister_parameter(self, param_name: str) -> None:
@@ -235,9 +258,12 @@ def set_weights(
         analog_tile = analog_tiles[0]
 
         if self.weight_scaling_omega > 0.0:
-            analog_tile.set_weights_scaled(weight, bias if self.analog_bias else None,
-                                           realistic=realistic,
-                                           omega=self.weight_scaling_omega)
+            analog_tile.set_weights_scaled(
+                weight, bias if self.analog_bias else None,
+                realistic=realistic,
+                omega=self.weight_scaling_omega,
+                weight_scaling_omega_columnwise=self.weight_scaling_omega_columnwise,
+                learn_out_scaling_alpha=self.learn_out_scaling_alpha)
         else:
             analog_tile.set_weights(weight, bias if self.analog_bias else None,
                                     realistic=realistic)
@@ -283,7 +309,9 @@ def get_weights(
 
         realistic = self.realistic_read_write and not force_exact
         if self.weight_scaling_omega > 0.0:
-            weight, bias = analog_tile.get_weights_scaled(realistic=realistic)
+            weight, bias = analog_tile.get_weights_scaled(
+                realistic=realistic,
+                weight_scaling_omega_columnwise=self.weight_scaling_omega_columnwise)
         else:
             weight, bias = analog_tile.get_weights(realistic=realistic)