Repository: IBM/aihwkit (language: python)

Issue:
Add experiments module
## Description and motivation

As a convenient way for running specific use cases (experiments) and encapsulating in an high-level and easy to handle and serialize construct, it would be good to add a new module. This new abstraction would facilitate exploring the toolkit and repeating similar use cases, and also facilitate future integrations.

## Proposed solution

Add an initial module with a specific experiment (training, based on the examples) and an interface for creating the experiment and executing it.

## Alternatives and other information

There are existing packages that provide support for "high level" pytorch interfacing. As the feature evolves, we should evaluate them and look into the possibility of making use of a suitable one - for this initial effort, the feature is envisaged to be small enough that a custom solution would fit the bill.

Add Experiments module with BasicTraining
## Related issues

#169 

## Description

First pass at #169, introducing a new `aihwkit.experiments` module containing a `BasicTraining` experiment, which represents training a neural network using a basic training loop.

## Details

<!-- A more elaborate description of the changes, if needed. -->



Discussion / hints:

Fix patch:
diff --git a/.pylintrc b/.pylintrc
index acfd7dd8..0464ed54 100644
--- a/.pylintrc
+++ b/.pylintrc
@@ -347,7 +347,7 @@ contextmanager-decorators=contextlib.contextmanager
 # List of members which are set dynamically and missed by pylint inference
 # system, and so shouldn't trigger E1101 when accessed. Python regular
 # expressions are accepted.
-generated-members=get_layer,get_tile,get_rpu_config,bias,preset_cls
+generated-members=get_layer,get_tile,get_rpu_config,bias,get_model,get_experiment,preset_cls
 
 # Tells whether missing members accessed in mixin class should be ignored. A
 # mixin class is detected if its name ends with "mixin" (case insensitive).
diff --git a/docs/source/api/aihwkit.experiments.rst b/docs/source/api/aihwkit.experiments.rst
index 6ef9cfb8..fad9ec13 100644
--- a/docs/source/api/aihwkit.experiments.rst
+++ b/docs/source/api/aihwkit.experiments.rst
@@ -13,3 +13,4 @@ Submodules
    :maxdepth: 4
 
    aihwkit.experiments.experiments
+   aihwkit.experiments.runners
diff --git a/docs/source/api/aihwkit.experiments.runners.base.rst b/docs/source/api/aihwkit.experiments.runners.base.rst
new file mode 100644
index 00000000..211b083c
--- /dev/null
+++ b/docs/source/api/aihwkit.experiments.runners.base.rst
@@ -0,0 +1,7 @@
+aihwkit.experiments.runners.base module
+=======================================
+
+.. automodule:: aihwkit.experiments.runners.base
+   :members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/docs/source/api/aihwkit.experiments.runners.local.rst b/docs/source/api/aihwkit.experiments.runners.local.rst
new file mode 100644
index 00000000..032f9d3b
--- /dev/null
+++ b/docs/source/api/aihwkit.experiments.runners.local.rst
@@ -0,0 +1,7 @@
+aihwkit.experiments.runners.local module
+========================================
+
+.. automodule:: aihwkit.experiments.runners.local
+   :members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/docs/source/api/aihwkit.experiments.runners.metrics.rst b/docs/source/api/aihwkit.experiments.runners.metrics.rst
new file mode 100644
index 00000000..0971c4e6
--- /dev/null
+++ b/docs/source/api/aihwkit.experiments.runners.metrics.rst
@@ -0,0 +1,7 @@
+aihwkit.experiments.runners.metrics module
+==========================================
+
+.. automodule:: aihwkit.experiments.runners.metrics
+   :members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/docs/source/api/aihwkit.experiments.runners.rst b/docs/source/api/aihwkit.experiments.runners.rst
new file mode 100644
index 00000000..5dfb9a6d
--- /dev/null
+++ b/docs/source/api/aihwkit.experiments.runners.rst
@@ -0,0 +1,17 @@
+aihwkit.experiments.runners package
+===================================
+
+.. automodule:: aihwkit.experiments.runners
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+Submodules
+----------
+
+.. toctree::
+   :maxdepth: 4
+
+   aihwkit.experiments.runners.base
+   aihwkit.experiments.runners.local
+   aihwkit.experiments.runners.metrics
diff --git a/docs/source/api_reference.rst b/docs/source/api_reference.rst
index 39142351..08c82e15 100644
--- a/docs/source/api_reference.rst
+++ b/docs/source/api_reference.rst
@@ -6,6 +6,9 @@ API Reference
 
    aihwkit
    aihwkit.exceptions
+   aihwkit.experiments
+   aihwkit.experiments.experiments
+   aihwkit.experiments.runners
    aihwkit.simulator
    aihwkit.simulator.configs
    aihwkit.simulator.noise_models
diff --git a/requirements.txt b/requirements.txt
index 2ba5b256..268ec45b 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -6,5 +6,6 @@ pybind11>=2.6.2
 torch>=1.5
 torchvision
 scipy
+requests>=2.25,<3
 numpy>=1.18
 dataclasses==0.7; python_version < '3.7'
diff --git a/setup.py b/setup.py
index d5c80718..e1f48cf4 100644
--- a/setup.py
+++ b/setup.py
@@ -19,6 +19,7 @@
     'torch{}'.format(os.getenv('TORCH_VERSION_SPECIFIER', '>=1.5')),
     'torchvision',
     'scipy',
+    'requests>=2.25,<3',
     'numpy>=1.18',
     'dataclasses==0.7; python_version < "3.7"'
 ]
diff --git a/src/aihwkit/experiments/runners/__init__.py b/src/aihwkit/experiments/runners/__init__.py
new file mode 100644
index 00000000..c6d16ae6
--- /dev/null
+++ b/src/aihwkit/experiments/runners/__init__.py
@@ -0,0 +1,17 @@
+# -*- 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.
+
+"""Experiment Runners for aihwkit."""
+
+# Convenience imports for easier access to the classes.
+
+from aihwkit.experiments.runners.local import LocalRunner
diff --git a/src/aihwkit/experiments/runners/base.py b/src/aihwkit/experiments/runners/base.py
new file mode 100644
index 00000000..a5e97d46
--- /dev/null
+++ b/src/aihwkit/experiments/runners/base.py
@@ -0,0 +1,27 @@
+# -*- 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.
+
+"""Base class for an Experiment Runner."""
+
+# pylint: disable=too-few-public-methods
+
+from typing import Any
+
+from aihwkit.experiments.experiments.base import Experiment
+
+
+class Runner:
+    """Base class for an Experiment Runner."""
+
+    def run(self, experiment: Experiment, **kwargs: Any) -> None:
+        """Run a single Experiment."""
+        raise NotImplementedError
diff --git a/src/aihwkit/experiments/runners/local.py b/src/aihwkit/experiments/runners/local.py
new file mode 100644
index 00000000..a9614cd9
--- /dev/null
+++ b/src/aihwkit/experiments/runners/local.py
@@ -0,0 +1,111 @@
+# -*- 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.
+
+"""Runner that executes Experiments locally."""
+
+from typing import Dict, List, Optional
+
+from torch import device as torch_device
+from torchvision.datasets import FashionMNIST, SVHN
+
+from aihwkit.experiments.experiments.base import Signals
+from aihwkit.experiments.runners.base import Runner
+from aihwkit.experiments.experiments.training import BasicTraining
+from aihwkit.experiments.runners.metrics import LocalMetric
+
+
+class LocalRunner(Runner):
+    """Runner that executes Experiments locally.
+
+    Class that allows executing Experiments locally.
+    """
+    # pylint: disable=too-few-public-methods
+
+    def __init__(self,
+                 device: Optional[torch_device] = None):
+        """LocalRunner constructor.
+
+        Args:
+            device: the device where the model will be moved to.
+        """
+        self.device = device
+
+    def run(  # type: ignore[override]
+            self,
+            experiment: BasicTraining,
+            max_elements_train: int = 0,
+            dataset_root: str = '/tmp/datasets',
+            stdout: bool = True
+    ) -> List[Dict]:
+        """Run a single Experiment.
+
+        Executes an experiment locally, in the device specified by
+        ``self.device``, optionally printing information to stdout.
+
+        Note:
+            If using a dataset different than ``FashionMNIST`` or ``SVHN``,
+            the runner assumes that the files for the dataset are downloaded at
+            ``dataset_root``. For those two datasets, the downloading will
+            take place automatically if the files are not present.
+
+        Args:
+            experiment: the experiment to be executed.
+            max_elements_train: limit on the amount of samples to use from
+                the dataset. If ``0``, no limit is applied.
+            dataset_root: path for the dataset files.
+            stdout: enable printing to stdout during the execution of the
+                experiment.
+
+        Returns:
+            A list of dictionaries with information about each epoch.
+        """
+        # pylint: disable=arguments-differ
+
+        # Setup the metric helper for the experiment.
+        metric = LocalMetric(stdout=stdout)
+        experiment.clear_hooks()
+        experiment.add_hook(Signals.EPOCH_START, metric.receive_epoch_start)
+        experiment.add_hook(Signals.EPOCH_END, metric.receive_epoch_end)
+        experiment.add_hook(Signals.TRAIN_EPOCH_END, metric.receive_train_epoch_end)
+        experiment.add_hook(Signals.TRAIN_EPOCH_BATCH_END, metric.receive_train_epoch_batch_end)
+        experiment.add_hook(Signals.VALIDATION_EPOCH_BATCH_END,
+                            metric.receive_validation_epoch_batch_end)
+        experiment.add_hook(Signals.VALIDATION_EPOCH_END, metric.receive_validation_epoch_end)
+
+        # Download the dataset if needed.
+        if experiment.dataset == FashionMNIST:
+            _ = experiment.dataset(dataset_root, download=True)
+        elif experiment.dataset == SVHN:
+            _ = experiment.dataset(dataset_root, download=True, split='train')
+            _ = experiment.dataset(dataset_root, download=True, split='test')
+
+        # Move the model to the device.
+        if self.device:
+            if self.device.type == 'cuda':
+                experiment.model = experiment.model.to(self.device)
+
+        # Build the objects needed for training.
+        training_loader, validation_loader = experiment.get_data_loaders(
+            experiment.dataset, experiment.batch_size,
+            max_elements_train=max_elements_train,
+            dataset_root=dataset_root
+        )
+        optimizer = experiment.get_optimizer(experiment.learning_rate, experiment.model)
+
+        experiment.results = experiment.train(training_loader,
+                                              validation_loader,
+                                              experiment.model,
+                                              optimizer,
+                                              experiment.loss_function(),
+                                              experiment.epochs,
+                                              self.device)
+        return experiment.results
diff --git a/src/aihwkit/experiments/runners/metrics.py b/src/aihwkit/experiments/runners/metrics.py
new file mode 100644
index 00000000..ea9abf3e
--- /dev/null
+++ b/src/aihwkit/experiments/runners/metrics.py
@@ -0,0 +1,101 @@
+# -*- 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.
+
+"""Helper for retrieving Metrics of an Experiment."""
+
+from datetime import datetime
+from typing import Dict, List
+
+
+class LocalMetric:
+    """Metric for local experiments.
+
+    Metric for local execution of experiments. Output to stdout can be
+    controlled by the ``stdout`` parameter to the constructor.
+    """
+
+    def __init__(self, stdout: bool = True) -> None:
+        self.epochs: List[Dict] = []
+        self.current_epoch: Dict = {}
+        self.stdout = stdout
+
+    def receive_epoch_start(self, epoch: int) -> None:
+        """Hook for `EPOCH_START`."""
+        self.current_epoch = {
+            'number': epoch,
+            'start_time': datetime.utcnow(),
+            'total_loss': 0,
+            'batches': 0,
+            'validation_batches': 0,
+            'validation_images': 0,
+            'validation_correct': 0,
+            'validation_loss': 0
+        }
+
+    def receive_train_epoch_batch_end(self, train_loss: float) -> None:
+        """Hook for `TRAIN_EPOCH_START`."""
+        self.current_epoch['batches'] += 1
+        self.current_epoch['total_loss'] += train_loss
+
+    def receive_validation_epoch_batch_end(
+            self,
+            total: int,
+            correct: int,
+            validation_loss: float
+    ) -> None:
+        """Hook for `VALIDATION_EPOCH_BATCH_END`."""
+        self.current_epoch['validation_images'] += total
+        self.current_epoch['validation_correct'] += int(correct)
+
+        self.current_epoch['validation_batches'] += 1
+        self.current_epoch['validation_loss'] += validation_loss
+
+    def receive_train_epoch_end(self) -> None:
+        """Hook for `TRAIN_EPOCH_END`."""
+        if not self.stdout:
+            return
+
+        print('Epoch: {}, loss: {:.8f}'.format(
+                self.current_epoch['number'],
+                self.current_epoch['total_loss'] / self.current_epoch['batches'],
+              ))
+
+    def receive_validation_epoch_end(self) -> None:
+        """Hook for `VALIDATION_EPOCH_END`."""
+        if not self.stdout:
+            return
+
+        print('Number of images: {}, accuracy: {:.6%}, validation loss: {:.8f}'.format(
+            self.current_epoch['validation_images'],
+            self.current_epoch['validation_correct'] / self.current_epoch['validation_images'],
+            self.current_epoch['validation_loss'] / self.current_epoch['batches'],
+        ))
+
+    def receive_epoch_end(self) -> Dict:
+        """Hook for `EPOCH_END`."""
+        end_time = datetime.utcnow()
+        time_epoch = (end_time - self.current_epoch['start_time']).total_seconds()
+        if self.stdout:
+            print('Time for epoch {}: {:.4}s'.format(self.current_epoch['number'],
+                                                     time_epoch))
+        self.epochs.append(self.current_epoch)
+
+        return {
+            'epoch': self.current_epoch['number'],
+            'time_epoch': time_epoch,
+            'accuracy': (self.current_epoch['validation_correct'] /
+                         self.current_epoch['validation_images']),
+            'train_loss': (self.current_epoch['total_loss'] /
+                           self.current_epoch['batches']),
+            'valid_loss': (self.current_epoch['validation_loss'] /
+                           self.current_epoch['batches'])
+        }


Test patch:
diff --git a/tests/helpers/decorators.py b/tests/helpers/decorators.py
index e931fb5b..4156bb42 100644
--- a/tests/helpers/decorators.py
+++ b/tests/helpers/decorators.py
@@ -102,3 +102,30 @@ def class_name(cls, _, params_dict):
     return parameterized_class(
         [{'preset_cls': preset} for preset in presets],
         class_name_func=class_name)
+
+
+def parametrize_over_models(models: List) -> Callable:
+    """Parametrize a TestCase over different kind of experiments.
+
+    Args:
+        models: list of model descriptions. The ``TestCase`` will be
+            repeated for each of the entries in the list.
+
+    Returns:
+        The decorated TestCase.
+    """
+
+    def object_to_dict(obj):
+        """Convert the public members of an object to a dictionary."""
+        ret = {key: value for key, value in vars(obj).items()
+               if not key.startswith('_')}
+        ret['parameter'] = obj.__name__
+
+        return ret
+
+    def class_name(cls, _, params_dict):
+        """Return a user-friendly name for a parametrized test."""
+        return '{}_{}'.format(cls.__name__, params_dict['parameter'])
+
+    return parameterized_class([object_to_dict(model) for model in models],
+                               class_name_func=class_name)
diff --git a/tests/helpers/experiments.py b/tests/helpers/experiments.py
new file mode 100644
index 00000000..79d2023f
--- /dev/null
+++ b/tests/helpers/experiments.py
@@ -0,0 +1,187 @@
+# -*- 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.
+
+# pylint: disable=missing-function-docstring,too-few-public-methods
+
+"""Models helpers for aihwkit tests."""
+
+from typing import Any
+
+from torch.nn import (
+    Flatten, LogSoftmax, MaxPool2d, Module, Sigmoid, Tanh
+)
+from torchvision.datasets import FashionMNIST, SVHN
+
+from aihwkit.experiments import BasicTraining
+from aihwkit.nn import AnalogConv2d, AnalogLinear, AnalogSequential
+from aihwkit.simulator.presets import (
+    CapacitorPreset, ReRamSBPreset, TikiTakaReRamSBPreset,
+    EcRamPreset
+)
+
+
+class FullyConnectedFashionMNIST:
+    """3 fully-connected layers; with FashionMNIST."""
+
+    def get_experiment(
+            self,
+            real: bool = False,
+            rpu_config: Any = CapacitorPreset
+    ):
+        """Return a BasicTraining experiment."""
+        argv = {
+            'dataset': FashionMNIST,
+            'model': self.get_model(rpu_config),
+            'epochs': 30,
+            'batch_size': 8,
+            'learning_rate': 0.01
+        }
+
+        if not real:
+            argv['epochs'] = 1
+
+        return BasicTraining(**argv)
+
+    def get_model(self, rpu_config: Any = CapacitorPreset) -> Module:
+        return AnalogSequential(
+            Flatten(),
+            AnalogLinear(784, 256, bias=True, rpu_config=rpu_config()),
+            Sigmoid(),
+            AnalogLinear(256, 128, bias=True, rpu_config=rpu_config()),
+            Sigmoid(),
+            AnalogLinear(128, 10, bias=True, rpu_config=rpu_config()),
+            LogSoftmax(dim=1)
+        )
+
+
+class FullyConnectedFashionMNISTTikiTaka:
+    """3 fully-connected layers; with FashionMNIST and tiki-taka."""
+
+    def get_experiment(
+            self,
+            real: bool = False,
+            rpu_config: Any = TikiTakaReRamSBPreset
+    ):
+        """Return a BasicTraining experiment."""
+        argv = {
+            'dataset': FashionMNIST,
+            'model': self.get_model(rpu_config),
+            'epochs': 30,
+            'batch_size': 8,
+            'learning_rate': 0.01
+        }
+
+        if not real:
+            argv['epochs'] = 1
+
+        return BasicTraining(**argv)
+
+    def get_model(self, rpu_config: Any = TikiTakaReRamSBPreset) -> Module:
+        return AnalogSequential(
+            Flatten(),
+            AnalogLinear(784, 256, bias=True, rpu_config=rpu_config()),
+            Sigmoid(),
+            AnalogLinear(256, 128, bias=True, rpu_config=rpu_config()),
+            Sigmoid(),
+            AnalogLinear(128, 10, bias=True, rpu_config=rpu_config()),
+            LogSoftmax(dim=1)
+        )
+
+
+class LeNet5FashionMNIST:
+    """LeNet5; with FashionMNIST."""
+
+    def get_experiment(
+            self,
+            real: bool = False,
+            rpu_config: Any = EcRamPreset
+    ):
+        """Return a BasicTraining experiment."""
+        argv = {
+            'dataset': FashionMNIST,
+            'model': self.get_model(rpu_config),
+            'epochs': 30,
+            'batch_size': 8,
+            'learning_rate': 0.01
+        }
+
+        if not real:
+            argv['epochs'] = 1
+
+        return BasicTraining(**argv)
+
+    def get_model(self, rpu_config: Any = EcRamPreset) -> Module:
+        return AnalogSequential(
+            AnalogConv2d(in_channels=1, out_channels=16, kernel_size=5, stride=1,
+                         rpu_config=rpu_config(),
+                         weight_scaling_omega=0.6),
+            Tanh(),
+            MaxPool2d(kernel_size=2),
+            AnalogConv2d(in_channels=16, out_channels=32, kernel_size=5, stride=1,
+                         rpu_config=rpu_config(),
+                         weight_scaling_omega=0.6),
+            Tanh(),
+            MaxPool2d(kernel_size=2),
+            Tanh(),
+            Flatten(),
+            AnalogLinear(in_features=512, out_features=128, rpu_config=rpu_config(),
+                         weight_scaling_omega=0.6),
+            Tanh(),
+            AnalogLinear(in_features=128, out_features=10, rpu_config=rpu_config(),
+                         weight_scaling_omega=0.6),
+            LogSoftmax(dim=1)
+        )
+
+
+class LeNet5SVHN:
+    """LeNet5; with SVHN."""
+
+    def get_experiment(
+            self,
+            real: bool = False,
+            rpu_config: Any = ReRamSBPreset
+    ):
+        """Return a BasicTraining experiment."""
+        argv = {
+            'dataset': SVHN,
+            'model': self.get_model(rpu_config),
+            'epochs': 30,
+            'batch_size': 8,
+            'learning_rate': 0.01
+        }
+
+        if not real:
+            argv['epochs'] = 1
+
+        return BasicTraining(**argv)
+
+    def get_model(self, rpu_config: Any = ReRamSBPreset) -> Module:
+        return AnalogSequential(
+            AnalogConv2d(in_channels=3, out_channels=16, kernel_size=5, stride=1,
+                         rpu_config=rpu_config(),
+                         weight_scaling_omega=0.6),
+            Tanh(),
+            MaxPool2d(kernel_size=2),
+            AnalogConv2d(in_channels=16, out_channels=32, kernel_size=5, stride=1,
+                         rpu_config=rpu_config(),
+                         weight_scaling_omega=0.6),
+            Tanh(),
+            MaxPool2d(kernel_size=2),
+            Tanh(),
+            Flatten(),
+            AnalogLinear(in_features=800, out_features=128, rpu_config=rpu_config(),
+                         weight_scaling_omega=0.6),
+            Tanh(),
+            AnalogLinear(in_features=128, out_features=10, rpu_config=rpu_config(),
+                         weight_scaling_omega=0.6),
+            LogSoftmax(dim=1)
+        )
diff --git a/tests/test_experiment_runners.py b/tests/test_experiment_runners.py
new file mode 100644
index 00000000..5a12a34e
--- /dev/null
+++ b/tests/test_experiment_runners.py
@@ -0,0 +1,71 @@
+# -*- 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 Experiment Runners."""
+
+from io import StringIO
+from unittest import skipIf
+from unittest.mock import patch
+
+from torch import device as torch_device
+from aihwkit.experiments.runners.local import LocalRunner
+from aihwkit.simulator.rpu_base import cuda
+
+from .helpers.decorators import parametrize_over_models
+from .helpers.experiments import (
+    FullyConnectedFashionMNIST, FullyConnectedFashionMNISTTikiTaka,
+    LeNet5FashionMNIST,
+)
+from .helpers.testcases import AihwkitTestCase
+
+
+@parametrize_over_models([
+    FullyConnectedFashionMNIST, FullyConnectedFashionMNISTTikiTaka,
+    LeNet5FashionMNIST,
+])
+class TestLocalRunner(AihwkitTestCase):
+    """Test LocalRunner."""
+
+    def test_run_example_cpu(self):
+        """Test running the example using a local runner."""
+        training_experiment = self.get_experiment()
+        local_runner = LocalRunner(device=torch_device('cpu'))
+
+        with patch('sys.stdout', new=StringIO()) as captured_stdout:
+            result = local_runner.run(training_experiment, max_elements_train=10)
+
+        # Asserts over stdout.
+        self.assertIn('Epoch: ', captured_stdout.getvalue())
+
+        # Asserts over the returned results.
+        self.assertEqual(len(result), 1)
+        self.assertEqual(result[0]['epoch'], 0)
+        self.assertIn('train_loss', result[0])
+        self.assertIn('accuracy', result[0])
+
+    @skipIf(not cuda.is_compiled(), 'not compiled with CUDA support')
+    def test_run_example_gpu(self):
+        """Test running the example using a local runner."""
+        training_experiment = self.get_experiment()
+        local_runner = LocalRunner(device=torch_device('cuda'))
+
+        with patch('sys.stdout', new=StringIO()) as captured_stdout:
+            result = local_runner.run(training_experiment, max_elements_train=10)
+
+        # Asserts over stdout.
+        self.assertIn('Epoch: ', captured_stdout.getvalue())
+
+        # Asserts over the returned results.
+        self.assertEqual(len(result), 1)
+        self.assertEqual(result[0]['epoch'], 0)
+        self.assertIn('train_loss', result[0])
+        self.assertIn('accuracy', result[0])