06_lenet5_hardware_aware.py

# -*- coding: utf-8 -*-

# (C) Copyright 2020, 2021, 2022 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 6: analog CNN with hardware aware training.

Mnist dataset on a LeNet5 inspired network.
"""
# pylint: disable=invalid-name

import os
from datetime import datetime

import matplotlib.pyplot as plt
import numpy as np

# Imports from PyTorch.
from torch import nn, device, manual_seed, no_grad
from torch import max as torch_max
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# Imports from aihwkit.
from aihwkit.nn import AnalogConv2d, AnalogLinear, AnalogSequential
from aihwkit.optim import AnalogSGD
from aihwkit.simulator.configs import (
    InferenceRPUConfig,
    WeightRemapType, WeightModifierType, WeightClipType,
    NoiseManagementType, BoundManagementType
)
from aihwkit.inference import PCMLikeNoiseModel
from aihwkit.simulator.rpu_base import cuda

# Check device
USE_CUDA = 0
if cuda.is_compiled():
    USE_CUDA = 1
DEVICE = device('cuda' if USE_CUDA else 'cpu')

# Path to store datasets
PATH_DATASET = os.path.join('data', 'DATASET')

# Path to store results
RESULTS = os.path.join(os.getcwd(), 'results', 'LENET5')
N_CLASSES = 10


def load_images(batch_size):
    """Load images for train from torchvision datasets.

    Args:
        batch_size (int): dtto

    Returns:
        DataLoader, DataLoader: train data and validation data
    """
    transform = transforms.Compose([transforms.ToTensor()])
    train_set = datasets.MNIST(PATH_DATASET, download=True, train=True, transform=transform)
    val_set = datasets.MNIST(PATH_DATASET, download=True, train=False, transform=transform)
    train_data = DataLoader(train_set, batch_size=batch_size, shuffle=True)
    validation_data = DataLoader(val_set, batch_size=batch_size, shuffle=False)

    return train_data, validation_data


def create_analog_network(rpu_config):
    """Return a LeNet5 inspired analog model.

    Args:
        rpu_config (InferenceRPUConfig): hardware and HWA training settings to use

    Returns:
        nn.Module: lenet analog model
    """
    channel = [16, 32, 512, 128]
    model = AnalogSequential(
        AnalogConv2d(in_channels=1, out_channels=channel[0], kernel_size=5, stride=1,
                     rpu_config=rpu_config),
        nn.Tanh(),
        nn.MaxPool2d(kernel_size=2),
        AnalogConv2d(in_channels=channel[0], out_channels=channel[1], kernel_size=5, stride=1,
                     rpu_config=rpu_config),
        nn.Tanh(),
        nn.MaxPool2d(kernel_size=2),
        nn.Tanh(),
        nn.Flatten(),
        AnalogLinear(in_features=channel[2], out_features=channel[3], rpu_config=rpu_config),
        nn.Tanh(),
        AnalogLinear(in_features=channel[3], out_features=N_CLASSES, rpu_config=rpu_config),
        nn.LogSoftmax(dim=1)
    )

    return model


def create_sgd_optimizer(model, learning_rate):
    """Create the analog-aware optimizer.

    Args:
        model (nn.Module): model to be trained
        learning_rate (float): global parameter to define learning rate

    Returns:
        Optimizer: created analog optimizer

    """
    optimizer = AnalogSGD(model.parameters(), lr=learning_rate)
    optimizer.regroup_param_groups(model)

    return optimizer


def train_step(data, model, criterion, optimizer):
    """Train network.

    Args:
        data (DataLoader): Validation set to perform the evaluation
        model (nn.Module): Trained model to be evaluated
        criterion (nn.CrossEntropyLoss): criterion to compute loss
        optimizer (Optimizer): analog model optimizer

    Returns:
        nn.Module, Optimizer, float: model, optimizer, and epoch loss
    """
    total_loss = 0

    model.train()

    for images, labels in data:
        images = images.to(DEVICE)
        labels = labels.to(DEVICE)
        optimizer.zero_grad()

        # Add training Tensor to the model (input).
        output = model(images)
        loss = criterion(output, labels)

        # Run training (backward propagation).
        loss.backward()

        # Optimize weights.
        optimizer.step()
        total_loss += loss.item() * images.size(0)
    epoch_loss = total_loss / len(data.dataset)

    return model, optimizer, epoch_loss


@no_grad()
def test_evaluation(data, model, criterion):
    """Test trained network

    Args:
        data (DataLoader): Validation set to perform the evaluation
        model (nn.Module): Trained model to be evaluated
        criterion (nn.CrossEntropyLoss): criterion to compute loss


    Returns:
        float, float, float: test epoch loss, test error, and test accuracy
    """
    total_loss = 0
    predicted_ok = 0
    total_images = 0

    model.eval()

    for images, labels in data:
        images = images.to(DEVICE)
        labels = labels.to(DEVICE)

        pred = model(images)
        loss = criterion(pred, labels)
        total_loss += loss.item() * images.size(0)

        _, predicted = torch_max(pred.data, 1)
        total_images += labels.size(0)
        predicted_ok += (predicted == labels).sum().item()

    accuracy = predicted_ok/total_images*100
    error = (1-predicted_ok/total_images)*100
    epoch_loss = total_loss / len(data.dataset)

    return epoch_loss, error, accuracy


def training_loop(model, criterion, optimizer, train_data, validation_data, epochs, print_every=1):
    """Training loop.

    Args:
        model (nn.Module): Trained model to be evaluated
        criterion (nn.CrossEntropyLoss): criterion to compute loss
        optimizer (Optimizer): analog model optimizer
        train_data (DataLoader): Validation set to perform the evaluation
        validation_data (DataLoader): Validation set to perform the evaluation
        epochs (int): global parameter to define epochs number
        print_every (int): defines how many times to print training progress

    Returns:
        nn.Module, Optimizer, Tuple: model, optimizer,
            and a tuple of lists of train losses, validation losses, and test
            error
    """
    train_losses = []
    valid_losses = []
    test_error = []

    # Train model
    for epoch in range(0, epochs):
        # Train_step
        model, optimizer, train_loss = train_step(train_data, model, criterion, optimizer)
        train_losses.append(train_loss)

        # Validate_step
        valid_loss, error, accuracy = test_evaluation(validation_data, model, criterion)
        valid_losses.append(valid_loss)
        test_error.append(error)

        if epoch % print_every == (print_every - 1):
            print(f'{datetime.now().time().replace(microsecond=0)} --- '
                  f'Epoch: {epoch}\t'
                  f'Train loss: {train_loss:.4f}\t'
                  f'Valid loss: {valid_loss:.4f}\t'
                  f'Test error: {error:.2f}%\t'
                  f'Accuracy: {accuracy:.2f}%\t')

    # Save results and plot figures
    np.savetxt(os.path.join(RESULTS, "Test_error.csv"), test_error, delimiter=",")
    np.savetxt(os.path.join(RESULTS, "Train_Losses.csv"), train_losses, delimiter=",")
    np.savetxt(os.path.join(RESULTS, "Valid_Losses.csv"), valid_losses, delimiter=",")

    return model, optimizer, (train_losses, valid_losses, test_error)


def plot_results(train_losses, valid_losses, test_error,
                 t_inference_times, inference_test_error):
    """Plot results.

    Args:
        train_losses (List): training losses as calculated in the training_loop
        valid_losses (List): validation losses as calculated in the training_loop
        test_error (List): test error as calculated in the training_loop
        t_inference_times (List): inference times
        inference_test_error (List): Inference test error
    """
    plt.ion()
    plt.figure(figsize=[14, 5])
    plt.subplot(1, 3, 1)
    h = plt.plot(train_losses, 'r-s', valid_losses, 'b-o')
    plt.title('LeNet5 - HWA training')
    plt.legend(h[:2], ['Training Losses', 'Validation Losses'])
    plt.xlabel('Epoch number')
    plt.ylabel('Loss [A.U.]')
    plt.grid(which='both', linestyle='--')

    plt.subplot(1, 3, 2)
    handle = plt.plot(test_error, 'r-s')
    plt.title('Test w/o prog. noise and drift')
    plt.legend(handle[:1], ['Validation test error'])
    plt.xlabel('Epoch number')
    plt.ylabel('Test Error [%]')
    plt.yscale('log')
    plt.ylim((5e-1, 1e2))
    plt.grid(which='both', linestyle='--')

    plt.subplot(1, 3, 3)
    handle = plt.plot(t_inference_times, inference_test_error, 'r-s')
    plt.title('Eval. w/ prog. noise and drift)')
    plt.legend(handle[:1], ['Validation test error'])
    plt.xlabel('Time of inference [s]')
    plt.ylabel('Test Error [%]')
    plt.yscale('log')
    plt.xscale('log')
    plt.ylim((5e-1, 1e2))
    plt.grid(which='both', linestyle='--')

    plt.show()
    plt.tight_layout()


def training_phase(model, criterion, optimizer, train_data, validation_data):
    """Training phase.

    Args:
        model (nn.Module): Trained model to be evaluated
        criterion (nn.CrossEntropyLoss): criterion to compute loss
        optimizer (Optimizer): analog model optimizer
        train_data (DataLoader): Validation set to perform the evaluation
        validation_data (DataLoader): Validation set to perform the evaluation

    Returns:
       Tuple: results from the training phase
    """
    print('\n ********************************************************* \n')
    print(f'\n{datetime.now().time().replace(microsecond=0)} --- '
          f'Started LeNet5 Training')

    model, optimizer, res = training_loop(model, criterion,
                                          optimizer, train_data,
                                          validation_data,
                                          N_EPOCHS)

    print(f'{datetime.now().time().replace(microsecond=0)} --- '
          f'Completed LeNet5 Training')

    return res


@no_grad()
def inference_phase(t_inference_times, model, criterion, validation_data):
    """Inference phase.

    Args:
        t_inference_times (list): list of times to do inference
        model (nn.Module): Trained model to be evaluated
        criterion (nn.CrossEntropyLoss): criterion to compute loss
        validation_data (DataLoader): Validation set to perform the evaluation

    Returns:
       Tuple: results from the training phase
    """
    # pylint: disable=too-many-locals

    _, error_pre, accuracy_pre = test_evaluation(validation_data, model, criterion)
    print(f'{datetime.now().time().replace(microsecond=0)} --- '
          f'Error after training: {error_pre:.2f}%\t'
          f'Accuracy after training: {accuracy_pre:.2f}%\t')

    error_lst = []
    accuracy_lst = []

    # Simulation of inference pass at different times after training.
    for t_inference in t_inference_times:
        model.drift_analog_weights(t_inference)

        _, error_post, accuracy_post = test_evaluation(validation_data, model, criterion)

        print(f'{datetime.now().time().replace(microsecond=0)} --- '
              f'Error after inference: {error_post:.2f}%\t'
              f'Accuracy after inference: {accuracy_post:.2f}%\t'
              f'Drift t={t_inference: .2e}\t')

        error_lst.append(error_post)
        accuracy_lst.append(accuracy_post)

    return error_lst, accuracy_lst


# Make sure the directory where to save the results exist.
# Results include: Loss vs Epoch graph, Accuracy vs Epoch graph and vector data.
os.makedirs(RESULTS, exist_ok=True)
manual_seed(1)

# Training parameters
N_EPOCHS = 30
BATCH_SIZE = 50
LEARNING_RATE = 0.1

# Load datasets.
training_data, valid_data = load_images(BATCH_SIZE)

# Define the properties of the neural network in terms of noise simulated during
# the inference/training pass
my_rpu_config = InferenceRPUConfig()
my_rpu_config.mapping.digital_bias = True
my_rpu_config.mapping.out_scaling_columnwise = True
my_rpu_config.mapping.learn_out_scaling = True
my_rpu_config.mapping.weight_scaling_omega = 1.0
my_rpu_config.mapping.weight_scaling_columnwise = False

my_rpu_config.noise_model = PCMLikeNoiseModel(g_max=25.0)
my_rpu_config.remap.type = WeightRemapType.CHANNELWISE_SYMMETRIC
my_rpu_config.clip.type = WeightClipType.LAYER_GAUSSIAN
my_rpu_config.clip.sigma = 2.5

# train input clipping
my_rpu_config.forward.noise_management = NoiseManagementType.NONE
my_rpu_config.forward.bound_management = BoundManagementType.NONE
my_rpu_config.forward.out_bound = 10.0  # quite restrictive
my_rpu_config.pre_post.input_range.enable = True
my_rpu_config.pre_post.input_range.manage_output_clipping = True
my_rpu_config.pre_post.input_range.decay = 0.001
my_rpu_config.pre_post.input_range.input_min_percentage = 0.95
my_rpu_config.pre_post.input_range.output_min_percentage = 0.95

my_rpu_config.modifier.type = WeightModifierType.ADD_NORMAL
my_rpu_config.modifier.std_dev = 0.1

# Prepare the model.
analog_model = create_analog_network(my_rpu_config)
if USE_CUDA:
    analog_model = analog_model.cuda()
print(analog_model)

opt = create_sgd_optimizer(analog_model, LEARNING_RATE)
crit = nn.CrossEntropyLoss()

# Train the model
results = training_phase(analog_model, crit, opt, training_data, valid_data)


# Test model inference over time
t_inference_lst = [0., 1., 20., 1000., 1e5, 1e7]
inference_error, _ = inference_phase(t_inference_lst, analog_model, crit, valid_data)
plot_results(*results, t_inference_lst, inference_error)