Source code for akida.model

import json
import numpy as np
import re

from .core import (Model, Layer, LayerType, LayerParams, AkidaUnsupervised, __version__,
                   evaluate_bitwidth)
from .statistics import Statistics
from .mapping import MapMode, _map_search


def model_str(self):
    data = "akida.Model, layer_count=" + str(self.layer_count)
    data += ", sequence_count=" + str(len(self.sequences))
    out_dims = self.output_shape if self.layer_count else []
    data += ", output_shape=" + str(out_dims)
    return data


def model_repr(self):
    out_dims = self.output_shape if self.layer_count else []
    data = "<akida.Model, layer_count=" + str(self.layer_count)
    data += ", output_shape=" + str(out_dims)
    data += ", sequences=" + repr(self.sequences) + ">"
    return data


def model_to_dict(self):
    """Provide a dict representation of the Model

    Returns:
        dict: a Model dictionary.
    """
    learning = None
    if self.learning:
        learning = {name: getattr(self.learning, name) for name in dir(self.learning)}
    return {
        "Akida version": __version__,
        "IP version": self.ip_version.name,
        "layers": [layer.to_dict() for layer in self.layers],
        "learning": learning,
        "input_shape": self.input_shape,
        "output_shape": self.output_shape
    }


def model_from_dict(model_dict):
    """Instantiate a Model from a dict representation

    Args:
        model_dict(dict): a Model dictionary.

    Returns:
        :obj:`Model`: a Model.
    """
    # Check major and minor version
    model_version = model_dict["Akida version"]
    major, minor, _ = __version__.split('.')
    model_major, model_minor, _ = model_version.split('.')
    if major != model_major or minor != model_minor:
        raise ValueError(f"Serialized model was generated by version {model_version}, which"
                         f" is incompatible with current version {__version__}")
    # Instantiate an empty Model
    model = Model()
    # Add layers
    layers = model_dict["layers"]
    for layer_dict in layers:
        layer_name = layer_dict["name"]
        # Extract layer parameters
        layer_params = layer_dict["parameters"]
        if layer_params is None:
            raise ValueError("Cannot deserialize a Layer without parameters")
        layer_params = layer_params.copy()
        # Evaluate the Layer type from its serialized name
        layer_type = getattr(LayerType, layer_params["layer_type"])
        # Remove layer_type
        layer_params.pop("layer_type")
        # Instantiate layer
        layer = Layer(LayerParams(layer_type, layer_params), layer_name)
        # Evaluate the layer inbounds
        inbounds = [model.get_layer(ib) for ib in layer_dict["inbounds"]]
        # Add it to the model
        model.add(layer, inbounds)
    # If needed compile it to set learning parameters and variables
    if model_dict["learning"] is not None:
        learning = model_dict["learning"].copy()
        # Some parameters must be integer
        for name in ["num_weights", "num_classes"]:
            learning[name] = int(learning[name])
        model.compile(AkidaUnsupervised(**learning))
    # Now that the Model is fully initialized, load layer variables
    for layer_dict in layers:
        layer_name = layer_dict["name"]
        # Get corresponding Layer in model
        layer = model.get_layer(layer_name)
        # Iterate over variables
        variables_dict = layer_dict["variables"]
        for name in variables_dict:
            variable_dict = variables_dict[name]
            dtype = variable_dict["dtype"]
            variable = np.array(variable_dict["data"]).astype(dtype)
            bitwidth = variable_dict["bitwidth"]
            actual_bitwidth = evaluate_bitwidth(variable)
            if bitwidth < actual_bitwidth:
                raise ValueError(f"The specified bitwidth ({bitwidth}) must be higher or equal to \
                                 the actual bitwidth ({actual_bitwidth}) of {name} \
                                 from the layer '{layer_name}'.")
            layer.variables[name] = variable
    return model


def model_to_json(self):
    """Provide a JSON representation of the Model

    Returns:
        str: a JSON-formatted string corresponding to a Model.
    """
    # Pretty-print serialized model
    model_str = json.dumps(self.to_dict(), indent=2)

    # Remove spurious line jumps in serialized arrays of numbers
    def align_arrays(m):
        # Just return the extracted pattern
        return m.group(1)
    # Look for lines starting with whitespaces and:
    # - a signed integer or float number with an optional opening square bracket,
    # - stand-alone opening or closing square brackets
    return re.sub(r'\n\s+(\[?-?[\d\.]+,?|\]|\[)', align_arrays, model_str)


def model_from_json(model_str):
    """Instantiate a Model from a JSON representation

    Args:
        model_str(str): a JSON-formatted string corresponding to a Model.

    Returns:
        :obj:`Model`: a Model.
    """
    return Model.from_dict(json.loads(model_str))


@property
def statistics(self):
    """Get statistics by sequence for this model.

    Returns:
       :obj:`Statistics`: the last inference statistics.

    """
    return Statistics(model=self)


def summary(self):
    """Prints a string summary of the model.

    This method prints a summary of the model with details for every layer,
    grouped by sequences:

    - name and type in the first column
    - output shape
    - kernel shape

    If there is any layer with unsupervised learning enabled, it will list
    them, with these details:

    - name of layer
    - number of incoming connections
    - number of weights per neuron

    """
    hrc_layers = [LayerType.InputConvolutional, LayerType.InputConv2D]

    def _model_summary(model):
        # prepare headers
        headers = ['Input shape', 'Output shape', 'Sequences', 'Layers']
        # prepare an empty table
        table = [headers]

        if not model.layers:
            row = [
                'N/A',
                'N/A',
                str(len(model.sequences)),
                str(len(model.layers))
            ]
        else:
            row = [
                str(model.input_shape),
                str(model.output_shape),
                str(len(model.sequences)),
                str(len(model.layers))
            ]
        # add the number of NPs if the model is mapped
        has_program = np.any([s.program is not None for s in self.sequences])
        if has_program:
            headers.append('NPs')
            headers.append('Skip DMAs')
            nb_nps = 0
            nb_skipdmas = 0
            for sequence in model.sequences:
                for current_pass in sequence.passes:
                    for layer in current_pass.layers:
                        if layer.parameters.layer_type not in hrc_layers:
                            nb_nps += 0 if layer.mapping is None else len(layer.mapping.nps)
                            nb_skipdmas += 0 if layer.mapping is None else len(
                                layer.mapping.skipdma_stores)
            row.append(nb_nps)
            row.append(nb_skipdmas)

            # External memory
            headers.append("External Memory (Bytes)")
            row.append(str(self.external_memory_size))

        table.append(row)
        print_table(table, "Model Summary")

    def _get_backend_info(sequence):
        backend = str(sequence.backend).split('.')[-1]
        sequence_info = sequence.name + " (" + backend + ")"
        if sequence.program is not None:
            sequence_info += " - size: " + str(len(sequence.program)) + " bytes"
        return sequence_info

    def _components_summary(model):
        headers = ['Component (type)', 'Count']
        # prepare an empty table
        table = [headers]
        component_count = model.component_count

        for comp_type, count in component_count.items():
            row = [comp_type.name, count]
            table.append(row)
        print_table(table, None)

    def _layer_components_summary(components, row):
        if components:
            assert all(c.type == components[0].type for c in components[1:])
            if row:
                row += " + "
            row += f"{len(components)} {components[0].type.name}"
        return row

    def _layers_summary(sequences):
        # Prepare headers
        headers = ['Layer (type)', 'Output shape', 'Kernel shape']
        has_program = np.any([s.program is not None for s in self.sequences])
        if has_program:
            headers.append('Components')
        # prepare an empty table
        table = [headers]
        new_splits = []
        for s in sequences:
            info = _get_backend_info(s)
            new_splits.append(info)
            nb_pass = 1
            for p in s.passes:
                if nb_pass > 1:
                    new_splits.append(f"pass {nb_pass}")
                for layer in p.layers:
                    nps = None if layer.mapping is None else layer.mapping.nps
                    skipdma_stores = None if layer.mapping is None else layer.mapping.skipdma_stores
                    skipdma_loads = None if layer.mapping is None else layer.mapping.skipdma_loads
                    # layer name (type)
                    layer_type = layer.parameters.layer_type
                    # kernel shape
                    if "weights" in layer.get_variable_names():
                        kernel_shape = layer.get_variable("weights").shape
                    elif layer_type == LayerType.StatefulRecurrent:
                        kernel_shape = str(layer.get_variable("in_proj").shape)
                        kernel_shape += "  "
                        kernel_shape += str(layer.get_variable("A_real").shape)
                        kernel_shape += "  "
                        kernel_shape += str(layer.get_variable("A_imag").shape)
                        kernel_shape += "  "
                        kernel_shape += str(layer.get_variable("out_proj").shape)
                    else:
                        kernel_shape = "N/A"
                    # Prepare row and add it
                    row = [str(layer), str(layer.output_dims), str(kernel_shape)]
                    if has_program:
                        components_list = ''
                        components_list = _layer_components_summary(nps, components_list)
                        components_list = _layer_components_summary(skipdma_stores, components_list)
                        components_list = _layer_components_summary(skipdma_loads, components_list)
                        row.append(components_list or "N/A")

                    table.append(row)
                    if len(table) - 1 > len(new_splits):
                        new_splits.append(False)
                    if layer_type == LayerType.SeparableConvolutional:
                        # Add pointwise weights on a second line
                        kernel_pw_shape = layer.get_variable("weights_pw").shape
                        row = ['', '', kernel_pw_shape]
                        if has_program:
                            row.append('')
                        table.append(row)
                        new_splits.append(False)
                nb_pass += 1
        print_table(table, None, new_splits)

    def _external_memory_summary(self):
        # Prepare headers
        headers = ['Layer (type)', 'External Memory (Bytes)']
        # prepare an empty table
        table = [headers]

        for s in self.sequences:
            for p in s.passes:
                for layer in p.layers:
                    if layer.mapping is None:
                        continue

                    # Fetch components
                    nps = layer.mapping.nps
                    skipdmas_store = layer.mapping.skipdma_stores

                    # Note: SkipDMA load does not use external memory; this is solely to
                    # maintain consistency with the loop over all components.
                    skipdmas_load = layer.mapping.skipdma_loads

                    layer_external_size = 0
                    for component in nps + skipdmas_store + skipdmas_load:
                        # For layer, we sum the external size over its components
                        layer_external_size += component.mem_info.external_size

                    if layer_external_size:
                        row = [str(layer), str(layer_external_size)]
                        table.append(row)

        print_table(table, "External Memory Summary")

    def _learning_summary(model):
        layer = model.layers[-1]
        # Prepare headers
        headers = ["Learning Layer", "# Input Conn.", "# Weights"]
        table = [headers]
        name = layer.name
        # Input connections is the product of input dims
        input_connections = np.prod(layer.input_dims)
        # Num non zero weights per neuron (counted on first neuron)
        weights = layer.get_variable("weights")
        incoming_conn = np.count_nonzero(weights[:, :, :, 0])
        # Prepare row and add it
        row = [name, str(input_connections), incoming_conn]
        table.append(row)
        print()
        print_table(table, "Learning Summary")

    # Print first the general Model summary
    _model_summary(self)
    # Print sequences summary
    if self.sequences:
        has_program = np.any([s.program is not None for s in self.sequences])
        if has_program:
            print()
            _components_summary(self)

            if self.external_memory_size:
                print()
                _external_memory_summary(self)

        print()
        _layers_summary(self.sequences)
    # Print learning summary if we have more than one input layer
    if len(self.layers) > 1 and self.learning:
        # Only the last layer of a model can learn
        print()
        _learning_summary(self)


def print_table(table, title, new_splits=None):
    # Convert to np.array
    to_str = np.vectorize(str, otypes=['O'])
    table = to_str(table)
    # get column lengths
    str_len_f = np.vectorize(lambda cell: len(str(cell)))
    str_lens = np.amax(str_len_f(table), 0)
    line_len = np.sum(str_lens)
    # Prepare format rows
    size_formats = np.vectorize(lambda cell: f"{{:{cell}.{cell}}}")
    format_strings = size_formats(str_lens)
    format_row = "  ".join(format_strings)
    # Generate separators
    separator_len = line_len + 2 * len(table[0])
    separator = "_" * separator_len
    double_separator = "=" * separator_len

    # Print header
    center_format = f"{{:^{separator_len}}}"
    if title is not None:
        print(center_format.format(title))
    print(separator)
    print(format_row.format(*table[0]))

    rows = table[1:, :]
    if new_splits is None:
        new_splits = [False] * len(rows)
    assert len(rows) == len(new_splits)
    if not any(new_splits):
        print(double_separator)
    # Print body
    for row, new_split in zip(rows, new_splits):
        if isinstance(new_split, str):
            # Display a line break only for sequences
            if "pass" not in new_split:
                print()
            # Compute the number of char on each side of the text
            space_len = max((separator_len - len(new_split)) / 2., 1.)
            space_left = "=" * int(np.ceil(space_len - 1))
            space_right = "=" * int(np.floor(space_len - 1))
            print(space_left, new_split, space_right)
            # Display a line break only for sequences
            if "pass" not in new_split:
                print()
        # Don't use a separator line on first row
        elif new_split is False and row[0] != rows[0][0]:
            print(separator)
        print(format_row.format(*row))
    print(separator)


def predict_classes(self, inputs, num_classes=0, batch_size=0):
    """Predicts the class labels for the specified inputs.

    Args:
        inputs (:obj:`numpy.ndarray`): a (n, x, y, c) uint8 tensor
        num_classes (int, optional): the number of output classes
        batch_size (int, optional): maximum number of inputs that should be
            processed at a time

    Returns:
        :obj:`numpy.ndarray`: an array of class labels

    """

    outputs = self.predict(inputs, batch_size)
    classes = np.argmax(outputs, axis=-1).flatten()
    num_neurons = outputs.shape[-1]
    if num_classes != 0 and num_classes != num_neurons:
        neurons_per_class = num_neurons // num_classes
        classes = classes // neurons_per_class
    return classes


def map(self, device, hw_only=False, mode=MapMode.AllNps, constraints=None):
    """Map the model to a Device using a target backend.

    This method tries to map a Model to the specified Device, implicitly
    identifying one or more layer sequences that are mapped individually on
    the Device Mesh.

    An optional hw_only parameter can be specified to force the mapping
    strategy to use only one hardware sequence, thus reducing software
    intervention on the inference.

    Note:
        Default mapping gives a higher throughput, lower latency, and better NP concurrent
        utilization but an optimal mapping depends on the system characteristics.

    Args:
        device (Device): the target Device or None
        hw_only (bool, optional): when true, the model should be mapped in one sequence.
            Defaults to False.
        mode (MapMode, optional): mapping mode. Defaults to MapMode.AllNps.
        constraints (MapConstraints, optional): mapping constraints. Defaults to None.
    """

    if mode == MapMode.Minimal or device is None:
        self._map(device, hw_only=hw_only, constraints=constraints)
    elif mode == MapMode.AllNps:
        _map_search(self, device, hw_only=hw_only, min_pass=True, constraints=constraints)
    elif mode == MapMode.HwPr:
        _map_search(self, device, hw_only=hw_only, min_pass=False, constraints=constraints)
    else:
        raise ValueError("Mapping Mode not supported.")