exportquant.py

import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import numpy as np
from datetime import datetime
from BitNetMCU import QuantizedModel
# from models import FCMNIST, CNNMNIST
import math
import matplotlib.pyplot as plt
import argparse
import yaml
import seaborn as sns
import importlib
from models import MaskingLayer

# Export quantized model from saved checkpoint
# cpldcpu 2024-04-14
# Note: Hyperparameters are used to generated the filename
#---------------------------------------------

showplots = False # display plots with statistics

def create_run_name(hyperparameters):
    runname = hyperparameters["runtag"] + '_' + hyperparameters["model"] + ('_Aug' if hyperparameters["augmentation"] else '') + '_BitMnist_' + hyperparameters["QuantType"] + "_width" + str(hyperparameters["network_width1"]) + "_" + str(hyperparameters["network_width2"]) + "_" + str(hyperparameters["network_width3"])  + "_epochs" + str(hyperparameters["num_epochs"])
    hyperparameters["runname"] = runname
    return runname

def load_model(model_name, params):
    """Instantiate model; forwards num_classes if provided."""
    try:
        module = importlib.import_module('models')
        model_class = getattr(module, model_name)
        kwargs = dict(
            network_width1=params["network_width1"],
            network_width2=params["network_width2"],
            network_width3=params["network_width3"],
            QuantType=params["QuantType"],
            NormType=params["NormType"],
            WScale=params["WScale"],
        )
        if 'cnn_width' in params:
            kwargs['cnn_width'] = params['cnn_width']
        if 'num_classes' in params:
            kwargs['num_classes'] = params['num_classes']
        return model_class(**kwargs)
    except AttributeError:
        raise ValueError(f"Model {model_name} not found in models.py")
    
def export_to_hfile(quantized_model, filename, runname, modelname=''):
    """
    Exports the quantized model to an Ansi-C header file.

    Parameters:
    filename (str): The name of the header file to which the quantized model will be exported.

    Note:
    This method currently only supports binary quantization.
    """

    if not quantized_model.quantized_model:
        raise ValueError("quantized_model is empty or None")

    # determine maximum number of activations per layer
    max_n_activations = max([layer['incoming_weights']  for layer in quantized_model.quantized_model if 'incoming_weights' in layer])
    # max_n_activations = 128

    with open(filename, 'w') as f:
        f.write(f'// Automatically generated header file\n')
        f.write(f'// Date: {datetime.now()}\n')
        f.write(f'// Quantized model exported from {runname}.pth\n')
        f.write('// Generated by exportquant.py\n\n')

        f.write('#include <stdint.h>\n\n')

        f.write('#ifndef BITNETMCU_MODEL_H\n')
        f.write('#define BITNETMCU_MODEL_H\n\n')

        f.write(f'// Model class name as defined in models.py\n')
        f.write(f'#define MODEL_{modelname}\n\n')

        f.write(f'// Number of layers\n')
        f.write(f'#define NUM_LAYERS {len(quantized_model.quantized_model)}\n\n')
        f.write(f'// Maximum number of activations per layer\n')
        f.write(f'#define MAX_N_ACTIVATIONS {max_n_activations}\n\n')

        for layer_info in quantized_model.quantized_model:
            layer = f'L{layer_info["layer_order"]}'

            if layer_info['layer_type'] == 'BitLinear':

                incoming_weights = layer_info['incoming_weights']
                outgoing_weights = layer_info['outgoing_weights']
                bpw = layer_info['bpw']
                weights = np.array(layer_info['quantized_weights'])
                quantization_type = layer_info['quantization_type']

                if (bpw*incoming_weights%32) != 0:
                    raise ValueError(f"Size mismatch: Incoming weights must be packed to 32bit boundary. Incoming weights: {incoming_weights} Bit per weight: {bpw} Total bits: {bpw*incoming_weights}")

                print(f'Layer: {layer} Quantization type: <{quantization_type}>, Bits per weight: {bpw}, Num. incoming: {incoming_weights},  Num outgoing: {outgoing_weights}')

                data_type = np.uint32

                if quantization_type == 'Binary':
                    encoded_weights = np.where(weights == -1, 0, 1)
                    QuantID = 1
                elif quantization_type == '2bitsym': # encoding -1.5 -> 11, -0.5 -> 10, 0.5 -> 00, 1.5 -> 01 (one complement with offset)
                    encoded_weights = ((weights < 0).astype(data_type) << 1) | (np.floor(np.abs(weights))).astype(data_type)  # use bitwise operations to encode the weights
                    QuantID = 2
                elif quantization_type == '4bitsym':
                    encoded_weights = ((weights < 0).astype(data_type) << 3) | (np.floor(np.abs(weights))).astype(data_type)  # use bitwise operations to encode the weights
                    QuantID = 4
                elif quantization_type == '4bit':
                    encoded_weights = np.floor(weights).astype(data_type) & 15  # twos complement encoding
                    QuantID =  8 + 4
                elif quantization_type == 'NF4':
                    levels = np.array([-1.0, -0.6962, -0.5251, -0.3949, -0.2844, -0.1848, -0.0911, 0.0,
                                   0.0796, 0.1609, 0.2461, 0.3379, 0.4407, 0.5626, 0.723, 1.0])
                    encoded_weights = np.argmin(np.abs(weights[:, :, np.newaxis] - levels), axis=2)
                    QuantID = 32 + 4  
                elif quantization_type == '8bit': 
                    encoded_weights = np.floor(weights).astype(data_type) & 255  # twos complement encoding
                    QuantID =  8 + 8
                elif quantization_type == 'FP130': # FP1.3.0 encoding (sign * 2^exp)
                    encoded_weights = ((weights < 0).astype(data_type) << 3) | (np.floor(np.log2(np.abs(weights)))).astype(data_type)
                    QuantID = 16 + 4
                elif quantization_type == 'Ternary':
                    # Ternary: 10 trits packed into 16 bits using base-3 encoding
                    # Encoding: +1=0, -1=1, 0=2
                    n_outputs, n_inputs = weights.shape

                    # Pad to multiple of 10 if needed
                    if n_inputs % 10 != 0:
                        pad_size = 10 - (n_inputs % 10)
                        print(f'WARNING: Ternary layer {layer} has {n_inputs} inputs, padding with {pad_size} zeros to align to 10')
                        weights = np.pad(weights, ((0, 0), (0, pad_size)), mode='constant', constant_values=0)
                        n_inputs = weights.shape[1]

                    # Map {-1, 0, 1} to {1, 2, 0} for efficient unpacking
                    trit_values = np.where(weights == 1, 0, np.where(weights == -1, 1, 2)).astype(np.uint32)

                    # Pack 10 trits into 16 bits per row
                    packed_row_size = n_inputs // 10
                    packed_weights = np.zeros((n_outputs, packed_row_size), dtype=np.uint16)

                    for row in range(n_outputs):
                        for word_idx in range(packed_row_size):
                            start = word_idx * 10
                            chunk = trit_values[row, start:start+10]

                            # Base-3 packing: MSB first
                            value = 0
                            for t in range(10):
                                value = value * 3 + chunk[t]
                            # Scale with ceiling: (value * 65536 + 59048) // 59049
                            packed = (value * 65536 + 59048) // 59049
                            packed_weights[row, word_idx] = packed

                    QuantID = 64  # Unique ID for ternary (0x40)

                    # Write ternary layer header (uint16_t arrays)
                    f.write(f'// Layer: {layer}\n')
                    f.write(f'// QuantType: {quantization_type} (10 trits per 16-bit word)\n')
                    f.write(f'#define {layer}_active\n')
                    f.write(f'#define {layer}_bitperweight {QuantID}\n')
                    f.write(f'#define {layer}_incoming_weights {n_inputs}\n')
                    f.write(f'#define {layer}_outgoing_weights {outgoing_weights}\n')

                    f.write(f'const uint16_t {layer}_weights[] = {{')
                    for i, data in enumerate(packed_weights.flatten()):
                        if i % 10 == 0:
                            f.write('\n\t')
                        f.write(f'0x{data:04x},')
                    f.write('\n};\n\n')

                    print(f'Layer: {layer} Ternary: {n_inputs} inputs (padded), {outgoing_weights} outputs, {packed_weights.size} uint16 words')
                    continue  # Skip standard 32-bit packing
                else:
                    print(f'Skipping layer {layer} with quantization type {quantization_type} and {bpw} bits per weight. Quantization type not supported.')
                    continue

                # pack bits into 32 bit words
                weight_per_word = 32 // bpw
                reshaped_array = encoded_weights.reshape(-1, weight_per_word)

                bit_positions = 32 - bpw - np.arange(weight_per_word, dtype=data_type) * bpw
                packed_weights = np.bitwise_or.reduce(reshaped_array << bit_positions, axis=1).view(data_type)

                # print(f'weights: {weights.shape} {weights.flatten()[0:16]}')
                # print(f'Encoded weights: {encoded_weights.shape} {encoded_weights.flatten()[0:16]}')
                # print(f'Packed weights: {packed_weights.shape} {", ".join(map(lambda x: hex(x), packed_weights.flatten()[0:4]))}')

                # Write layer order, shape, shiftright and weights to the file
                f.write(f'// Layer: {layer}\n')
                f.write(f'// QuantType: {quantization_type}\n')

                f.write(f'#define {layer}_active\n')
                f.write(f'#define {layer}_bitperweight {QuantID} \n')
                f.write(f'#define {layer}_incoming_weights {incoming_weights}\n')
                f.write(f'#define {layer}_outgoing_weights {outgoing_weights}\n')

                f.write(f'const uint32_t {layer}_weights[] = {{')
                for i,data in enumerate(packed_weights.flatten()):
                    if i&7 ==0:
                        f.write('\n\t')
                    f.write(f'0x{data:08x},')
                f.write('\n}; //first channel is topmost bit\n\n')

            elif layer_info['layer_type'] == 'BitConv2d':
                in_channels = layer_info['in_channels']
                out_channels = layer_info['out_channels']
                incoming_x = layer_info['incoming_x']
                incoming_y = layer_info['incoming_y']
                outgoing_x = layer_info['outgoing_x']
                outgoing_y = layer_info['outgoing_y']

                groups = layer_info['groups']
                kernel_size = layer_info['kernel_size'][0]  # Assuming square kernel
                bpw = layer_info['bpw']
                weights = np.array(layer_info['quantized_weights'])

                f.write(f'// Layer: {layer} (Convolutional)\n')
                f.write(f'#define {layer}_active\n')
                f.write(f'#define {layer}_type BitConv2d\n')
                f.write(f'#define {layer}_in_channels {in_channels}\n')
                f.write(f'#define {layer}_out_channels {out_channels}\n')
                f.write(f'#define {layer}_incoming_x {incoming_x}\n')
                f.write(f'#define {layer}_incoming_y {incoming_y}\n')
                f.write(f'#define {layer}_outgoing_x {outgoing_x}\n')
                f.write(f'#define {layer}_outgoing_y {outgoing_y}\n')
                f.write(f'#define {layer}_kernel_size {kernel_size}\n')
                f.write(f'#define {layer}_stride 1\n')
                f.write(f'#define {layer}_padding 0\n')
                f.write(f'#define {layer}_groups {groups}\n')
                f.write(f'#define {layer}_bitperweight {bpw}\n')

                f.write(f'const int8_t {layer}_weights[] = {{')
                for i, data in enumerate(weights.flatten()):
                    if i % 16 == 0:
                        f.write('\n\t')
                    f.write(f'{int(data)},')
                f.write('\n};\n\n')                

                print(f'Layer: {layer} Conv2d bpw: {bpw} {in_channels} -> {out_channels} groups:{groups} Kernel: {kernel_size}x{kernel_size} Incoming: {incoming_x}x{incoming_y} Outgoing: {outgoing_x}x{outgoing_y}')                

            elif layer_info['layer_type'] == 'MaxPool2d':
                pool_size = layer_info['kernel_size']
                incoming_x = layer_info['incoming_x']
                incoming_y = layer_info['incoming_y']
                outgoing_x = layer_info['outgoing_x']
                outgoing_y = layer_info['outgoing_y']

                f.write(f'#define {layer}_active\n')
                f.write(f'#define {layer}_type MaxPool2d\n')
                f.write(f'#define {layer}_pool_size {pool_size}\n')
                f.write(f'#define {layer}_incoming_x {incoming_x}\n')
                f.write(f'#define {layer}_incoming_y {incoming_y}\n')
                f.write(f'#define {layer}_outgoing_x {outgoing_x}\n')
                f.write(f'#define {layer}_outgoing_y {outgoing_y}\n\n')

                print(f'Layer: {layer} MaxPool2d Pool Size: {pool_size} Incoming: {incoming_x}x{incoming_y} Outgoing: {outgoing_x}x{outgoing_y}')
               
        f.write('#endif\n')

def plot_test_images(test_loader):
    dataiter = iter(test_loader)
    images, labels = next(dataiter)

    fig, axes = plt.subplots(5, 5, figsize=(8, 8))

    for i, ax in enumerate(axes.flat):
        ax.imshow(images[i].numpy().squeeze(), cmap='gray')
        ax.set_title(f'Label: {labels[i]}')
        ax.axis('off')

    plt.tight_layout()
    plt.show()

def print_stats(quantized_model):
    for layer_info in quantized_model.quantized_model:
        if layer_info['layer_type'] == 'BitLinear' or layer_info['layer_type'] == 'BitConv2d':
            weights = np.array(layer_info['quantized_weights'])
            print()
            print(f'Layer: {layer_info["layer_order"]}, Max: {np.max(weights)}, Min: {np.min(weights)}, Mean: {np.mean(weights)}, Std: {np.std(weights)}')

            values, counts = np.unique(weights, return_counts=True)
            probabilities = counts / np.sum(counts)

            print(f'Values: {values}')
            print(f'Percent: {(probabilities * 100)}')

            number_of_codes = 2**layer_info['bpw'] 
            entropy = -np.sum(probabilities * np.log2(probabilities))
            print(f'Entropy: {entropy:.2f} bits. Code capacity used: {entropy / np.log2(number_of_codes) * 100} %')
  

def plot_statistics(quantized_model):
    # Step 1: Extract the weights of the first layer
    first_layer_weights = np.array(quantized_model.quantized_model[0]['quantized_weights'])

    # Step 2: Reshape the weights into a 16x16 grid
    reshaped_weights = first_layer_weights.reshape(16, 16, -1)
    print(reshaped_weights.shape)
    # Step 3: Calculate the variance of each channel
    variances = np.var(reshaped_weights, axis=-1)

    # Calculate the mean of each channel
    means = np.mean(reshaped_weights, axis=-1)

    # Create a figure with 2 subplots: one for variance, one for mean
    fig, axs = plt.subplots(1, 2, figsize=(12, 6))

    # Plot the variance
    axs[0].imshow(variances, cmap='hot', interpolation='nearest')
    axs[0].set_title('Variance vs Channel')
    fig.colorbar(plt.cm.ScalarMappable(cmap='hot'), ax=axs[0], label='Variance')

    # Plot the mean
    im = axs[1].imshow(means, cmap='hot', interpolation='nearest')
    axs[1].set_title('Mean vs Channel')
    fig.colorbar(im, ax=axs[1], label='Mean')

    # Display the plot
    plt.show(block=False)


def plot_weights(quantized_model):
    # Step 1: Extract the weights of the first layer
    first_layer_weights = np.array(quantized_model.quantized_model[0]['quantized_weights'])

    # Step 2: Reshape the weights into a 16x16 grid for each output channel
    reshaped_weights = first_layer_weights.reshape(-1, 16, 16)

    # Calculate the number of output channels
    num_channels = reshaped_weights.shape[0]

    # Calculate the number of rows and columns for the subplots
    num_cols = int(math.sqrt(num_channels))
    num_rows = num_channels // num_cols
    if num_channels % num_cols != 0:
        num_rows += 1

    # Step 3: Create a figure with a grid of subplots, one for each output channel
    fig, axs = plt.subplots(num_rows, num_cols, figsize=(8, 8))

    # Step 4: For each output channel, plot the weights in the corresponding subplot
    for i in range(num_cols*num_rows):
        row = i // num_cols
        col = i % num_cols
        if i < num_channels:
            axs[row, col].imshow(reshaped_weights[i], cmap='hot', interpolation='nearest')
        axs[row, col].axis('off')  # Turn off axis for each subplot
    # Reduce the gaps between the subplots
    # plt.subplots_adjust(wspace=-0.10, hspace=-0.10)
    # Display the plot
    plt.tight_layout()  # This will ensure the subplots do not overlap
    plt.show(block=False)

def plot_weight_histograms(quantized_model):
    fig = plt.figure(figsize=(10, 10))

    for layer_index, layer in enumerate(quantized_model.quantized_model):
        if layer['layer_type'] == 'BitConv2d' or layer['layer_type'] == 'BitLinear':
            layer_weights = np.array(layer['quantized_weights'])
            bpw = layer['bpw']

            flattened_weights = layer_weights.flatten()

            ax = fig.add_subplot(len(quantized_model.quantized_model), 1, layer_index + 1)

            # ax.hist(flattened_weights, width=1, bins='auto')
            sns.histplot(flattened_weights, bins=2**bpw, ax=ax, kde=True)
            ax.set_title(f'Layer {layer_index+1} Weight Distribution')

    plt.tight_layout()
    plt.show(block=False)

def print_masking_layers(model):
    masking_layer_count = 0
    for name, module in model.named_modules():
        if isinstance(module, MaskingLayer):
            masking_layer_count += 1
            mask_params = module.mask.data
            # mask = torch.sigmoid(mask_params )
            mask = mask_params
            active_channels = torch.sum(mask > 0.5).item()
            total_channels = len(mask)
            active_percentage = (active_channels / total_channels) * 100
            print(mask_params)
            print(f"SoftMaskLayer found in {name}:")
            print(f"  Total channels: {total_channels}")
            print(f"  Active channels (mask > 0.5): {active_channels}")
            print(f"  Pruned channels (mask <= 0.5): {total_channels - active_channels}")
            print(f"  Active percentage: {active_percentage:.2f}%")
            
            print("  Mask parameter statistics:")
            print(f"    Min: {mask_params.min().item():.4f}")
            print(f"    Max: {mask_params.max().item():.4f}")
            print(f"    Mean: {mask_params.mean().item():.4f}")
            print(f"    Std: {mask_params.std().item():.4f}")

            print("  Soft mask value histogram:")
            hist = torch.histc(mask_params, bins=10)
            bin_edges = torch.linspace(mask_params.min(), mask_params.max(), steps=11)
            for i, count in enumerate(hist):
                print(f"    {bin_edges[i]:.1f}-{bin_edges[i+1]:.1f}: {count.item():.0f}")
            
            print("\n")

    if masking_layer_count == 0:
        print("No SoftMaskLayers found in the model.")
    else:
        print(f"Total SoftMaskLayers found: {masking_layer_count}")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Training script')
    parser.add_argument('--params', type=str, help='Name of the parameter file', default='trainingparameters.yaml')

    args = parser.parse_args()

    if args.params:
        paramname = args.params
    else:
        paramname = 'trainingparameters.yaml'

    print(f'Load parameters from file: {paramname}')
    with open(paramname) as f:
        hyperparameters = yaml.safe_load(f)

    # main
    runname= create_run_name(hyperparameters)
    print(runname)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # Dataset selection (mirror training.py logic)
    dataset_name = hyperparameters.get("dataset", "MNIST").upper()

    if dataset_name == "MNIST":
        num_classes = 10
        mean, std = (0.1307,), (0.3081,)
        base_dataset_train = datasets.MNIST
        base_dataset_test = datasets.MNIST
        dataset_kwargs = {"train": True}
        dataset_kwargs_test = {"train": False}
    elif dataset_name.startswith("EMNIST"):
        split = dataset_name.split('_')[1].lower() if '_' in dataset_name else 'balanced'
        split_alias = { 'BALANCED':'balanced', 'BYCLASS':'byclass', 'BYMERGE':'bymerge', 'LETTERS':'letters', 'DIGITS':'digits', 'MNIST':'mnist'}
        split = split_alias.get(split.upper(), split)
        split_classes = { 'byclass':62, 'bymerge':47, 'balanced':47, 'letters':37, 'digits':10, 'mnist':10 }
        num_classes = split_classes.get(split, 47)
        from torchvision.datasets import EMNIST
        mean, std = (0.1307,), (0.3081,)
        base_dataset_train = EMNIST
        base_dataset_test = EMNIST
        dataset_kwargs = {"split": split, "train": True}
        dataset_kwargs_test = {"split": split, "train": False}
    else:
        raise ValueError(f"Unsupported dataset: {dataset_name}")

    transform = transforms.Compose([
        transforms.Resize((16, 16)),
        transforms.ToTensor(),
        transforms.Normalize(mean, std)
    ])

    # Only test set strictly needed here, but keep parity
    train_data = base_dataset_train(root='data', transform=transform, download=True, **dataset_kwargs)
    test_data = base_dataset_test(root='data', transform=transform, download=True, **dataset_kwargs_test)

    hyperparameters['num_classes'] = num_classes
    # Create data loaders
    test_loader = DataLoader(test_data, batch_size=hyperparameters["batch_size"], shuffle=False)

    model = load_model(hyperparameters["model"], hyperparameters).to(device)

    print('Loading model...')
    checkpoint_path = f'modeldata/{runname}.pth'
    try:
        state = torch.load(checkpoint_path, map_location='cpu')
    except FileNotFoundError:
        print(f"The file '{checkpoint_path}' does not exist.")
        exit()

    # Attempt load; if classifier size mismatch, rebuild with inferred num_classes from checkpoint
    try:
        model.load_state_dict(state, strict=True)
    except RuntimeError as e:
        if 'classifier.weight' in str(e):
            saved_classes = state['classifier.weight'].shape[0]
            print(f"Classifier size mismatch. Rebuilding model with num_classes={saved_classes} (was {hyperparameters.get('num_classes')}).")
            hyperparameters['num_classes'] = saved_classes
            model = load_model(hyperparameters["model"], hyperparameters).to(device)
            model.load_state_dict(state, strict=True)
        else:
            raise

    print('Inference using the original model...')
    correct = 0
    total = 0
    test_loss = []
    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    testaccuracy = correct / total * 100
    print(f'Accuracy/Test of trained model: {testaccuracy} %')

    print('Quantizing model...')
    # Quantize the model
    quantized_model = QuantizedModel(model)

    # Print statistics
    print_stats(quantized_model)

    print_masking_layers(model)

    if showplots:
        # plot_weights(quantized_model)
        # plot_statistics(quantized_model)
        plot_weight_histograms(quantized_model)
        # plot_test_images(test_loader)

    print(f'Total number of bits: {quantized_model.totalbits()} ({quantized_model.totalbits()/8/1024} kbytes)')

    # Inference using the quantized model
    print ("inference of quantized model")

    # Initialize counters
    total_correct_predictions = 0
    total_samples = 0

    # Iterate over the test data
    for input_data, labels in test_loader:
        # Reshape and convert to numpy
        input_data = input_data.view(input_data.size(0), -1).cpu().numpy()

        labels = labels.cpu().numpy()

        # Inference
        result = quantized_model.inference_quantized(input_data)

        # Get predictions
        predict = np.argmax(result, axis=1)

        # Calculate the fraction of correct predictions for this batch
        correct_predictions = (predict == labels).sum()

        # Update counters
        total_correct_predictions += correct_predictions  # Multiply by batch size
        total_samples += input_data.shape[0]

    # Calculate and print the overall fraction of correct predictions
    overall_correct_predictions = total_correct_predictions / total_samples

    print('Accuracy/Test of quantized model:', overall_correct_predictions * 100, '%')

    print("Exporting model to header file")
    # export the quantized model to a header file
    # export_to_hfile(quantized_model, f'{exportfolder}/{runname}.h')
    export_to_hfile(quantized_model, f'BitNetMCU_model.h',runname, hyperparameters["model"])
    
    if showplots:
        plt.show()