From 7ee0f44c8ae055872dadc8ff5721f9f04a68d1b5 Mon Sep 17 00:00:00 2001 From: Pablo Carmona Gonzalez Date: Mon, 12 Aug 2024 13:08:45 +0200 Subject: [PATCH 1/3] feat(docs): add half-precision training section in using_simulator docs Signed-off-by: Pablo Carmona Gonzalez --- docs/source/using_simulator.rst | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/docs/source/using_simulator.rst b/docs/source/using_simulator.rst index 25688ba3..6b257470 100644 --- a/docs/source/using_simulator.rst +++ b/docs/source/using_simulator.rst @@ -399,6 +399,21 @@ instead of manually specifying a ``RPU Configuration``:: tile = AnalogTile(10, 20, rpu_config=TikiTakaEcRamPreset()) +Working with half-precision training +------------------------------------ + +The simulator supports half-precision training. This can be enabled by setting the +``RPUDataType`` to ``HALF`` when creating the configuration:: + + from aihwkit.simulator.configs import InferenceRPUConfig + from aihwkit.simulator.parameters.enums import RPUDataType + + rpu_config = InferenceRPUConfig() # or TorchInferenceRPUConfig(). + rpu_config.runtime.data_type = RPUDataType.HALF + +For more info look into :py:mod:`aihwkit.simulator.parameters.enums.RPUDataType`. + + .. _Gokmen & Haensch 2020: https://www.frontiersin.org/articles/10.3389/fnins.2020.00103/full .. _Example 7: https://github.com/IBM/aihwkit/blob/master/examples/07_simple_layer_with_other_devices.py .. _Example 8: https://github.com/IBM/aihwkit/blob/master/examples/08_simple_layer_with_tiki_taka.py From 0877c2363801154f962e06600709d8dcbe90f7f0 Mon Sep 17 00:00:00 2001 From: Pablo Carmona Gonzalez Date: Mon, 12 Aug 2024 17:53:53 +0200 Subject: [PATCH 2/3] feat(examples): add example 31 for half precision training Signed-off-by: Pablo Carmona Gonzalez --- examples/31_half_precision_training.py | 83 ++++++++++++++++++++++++++ 1 file changed, 83 insertions(+) create mode 100644 examples/31_half_precision_training.py diff --git a/examples/31_half_precision_training.py b/examples/31_half_precision_training.py new file mode 100644 index 00000000..ad4a57e4 --- /dev/null +++ b/examples/31_half_precision_training.py @@ -0,0 +1,83 @@ +# type: ignore +# pylint: disable-all +# -*- coding: utf-8 -*- + +# (C) Copyright 2020, 2021, 2022, 2023, 2024 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. + +"""aihwkit example 31: Using half precision training. + +This example demonstrates how to use half precision training with aihwkit. + +""" +# pylint: disable=invalid-name + +import tqdm +import torch +import torch.nn as nn +import torch.nn.functional as F +from torchvision import datasets, transforms +from aihwkit.simulator.configs import InferenceRPUConfig, TorchInferenceRPUConfig +from aihwkit.nn.conversion import convert_to_analog +from aihwkit.optim import AnalogSGD +from aihwkit.simulator.parameters.enums import RPUDataType + +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + +class Net(nn.Module): + def __init__(self): + super(Net, self).__init__() + self.conv1 = nn.Conv2d(1, 32, 3, 1) + self.conv2 = nn.Conv2d(32, 64, 3, 1) + self.dropout1 = nn.Dropout(0.25) + self.dropout2 = nn.Dropout(0.5) + self.fc1 = nn.Linear(9216, 128) + self.fc2 = nn.Linear(128, 10) + + def forward(self, x): + x = self.conv1(x) + x = F.relu(x) + x = self.conv2(x) + x = F.relu(x) + x = F.max_pool2d(x, 2) + x = self.dropout1(x) + x = torch.flatten(x, 1) + x = self.fc1(x) + x = F.relu(x) + x = self.dropout2(x) + x = self.fc2(x) + output = F.log_softmax(x, dim=1) + return output + +if __name__ == "__main__": + model = Net() + rpu_config = TorchInferenceRPUConfig() + model = convert_to_analog(model, rpu_config) + nll_loss = torch.nn.NLLLoss() + transform=transforms.Compose([ + transforms.ToTensor(), + transforms.Normalize((0.1307,), (0.3081,)) + ]) + dataset = datasets.MNIST('data', train=True, download=True, transform=transform) + train_loader = torch.utils.data.DataLoader(dataset, batch_size=32) + + model = model.to(device=device, dtype=torch.bfloat16) + optimizer = AnalogSGD(model.parameters(), lr=0.1) + model = model.train() + + pbar = tqdm.tqdm(enumerate(train_loader)) + for batch_idx, (data, target) in pbar: + data, target = data.to(device=device, dtype=torch.bfloat16), target.to(device=device) + optimizer.zero_grad() + output = model(data) + loss = F.nll_loss(output.float(), target) + loss.backward() + optimizer.step() + pbar.set_description(f"Loss {loss:.4f}") \ No newline at end of file From 4e1bf93c00d2bfc5389fb2c5daf3d0e911027585 Mon Sep 17 00:00:00 2001 From: Pablo Carmona Gonzalez Date: Mon, 19 Aug 2024 12:24:26 +0200 Subject: [PATCH 3/3] format example with black --- examples/31_half_precision_training.py | 25 ++++++++++++++----------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/examples/31_half_precision_training.py b/examples/31_half_precision_training.py index ad4a57e4..69ae9eb6 100644 --- a/examples/31_half_precision_training.py +++ b/examples/31_half_precision_training.py @@ -31,6 +31,7 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + class Net(nn.Module): def __init__(self): super(Net, self).__init__() @@ -55,29 +56,31 @@ def forward(self, x): x = self.fc2(x) output = F.log_softmax(x, dim=1) return output - + + if __name__ == "__main__": model = Net() rpu_config = TorchInferenceRPUConfig() - model = convert_to_analog(model, rpu_config) + model = convert_to_analog(model, rpu_config) nll_loss = torch.nn.NLLLoss() - transform=transforms.Compose([ - transforms.ToTensor(), - transforms.Normalize((0.1307,), (0.3081,)) - ]) - dataset = datasets.MNIST('data', train=True, download=True, transform=transform) + transform = transforms.Compose( + [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))] + ) + dataset = datasets.MNIST("data", train=True, download=True, transform=transform) train_loader = torch.utils.data.DataLoader(dataset, batch_size=32) - + model = model.to(device=device, dtype=torch.bfloat16) optimizer = AnalogSGD(model.parameters(), lr=0.1) model = model.train() - + pbar = tqdm.tqdm(enumerate(train_loader)) for batch_idx, (data, target) in pbar: - data, target = data.to(device=device, dtype=torch.bfloat16), target.to(device=device) + data, target = data.to(device=device, dtype=torch.bfloat16), target.to( + device=device + ) optimizer.zero_grad() output = model(data) loss = F.nll_loss(output.float(), target) loss.backward() optimizer.step() - pbar.set_description(f"Loss {loss:.4f}") \ No newline at end of file + pbar.set_description(f"Loss {loss:.4f}")