Extending nnio

Using our API to wrap around your own custom models

nnio.Model is an abstract class from which all models in nnio are derived. It is easy to use by redefining forward method:

class MyClassifier(nnio.Model):
    def __init__(self):
        super().__init__()
        self.model = SomeModel()

    def forward(self, image):
        # Do something with image
        result = self.model(image)
        # For example, classification
        if result == 0:
            return 'person'
        else:
            return 'cat'

    def get_preprocessing(self):
        return nnio.Preprocessing(
            resize=(224, 224),
            dtype='float',
            divide_by_255=True,
            means=[0.485, 0.456, 0.406],
            stds=[0.229, 0.224, 0.225],
            batch_dimension=True,
            channels_first=True,
        )

We also recommend to define get_preprocessing method like in Model Zoo models. See nnio.Preprocessing. We encourage users to wrap their loaded models in such classes. nnio.Model abstract base class is described below:

nnio.Model

class nnio.Model

abstract forward(*args, **kwargs)

This method is called when the model is called.

Parameters

• *inputs – numpy arrays, Inputs to the model

• return_info – bool, If True, will return inference time

Returns

numpy array or list of numpy arrays.

get_input_details()

Returns

human-readable model input details.

get_output_details()

Returns

human-readable model output details.

get_preprocessing()

Returns

nnio.Preprocessing object.