Modified ML model to ONNX format

pytest-benchmark

numpy

scikit-learn

onnxruntime

import os

import grpc

import numpy as np

import pickle as pkl

import onnxruntime as rt

import logging

from l3_centralizedattackdetector.proto.l3_centralizedattackdetector_pb2 import (

    Empty,

LOGGER = logging.getLogger(__name__)

here = os.path.dirname(os.path.abspath(__file__))

MODEL_FILE = os.path.join(here, "ml_model/RF_Netflow_TF")

MODEL_FILE = os.path.join(here, "ml_model/teraflow_rf.onnx")

class l3_centralizedattackdetectorServiceServicerImpl(L3CentralizedattackdetectorServicer):

    def __init__(self):

        LOGGER.debug("Creating Servicer...")

        self.inference_values = []

        with open(MODEL_FILE, "rb") as f:

            self.ml_model = pkl.load(f)

        self.model = rt.InferenceSession(MODEL_FILE)

        self.input_name = self.model.get_inputs()[0].name

        self.label_name = self.model.get_outputs()[0].name

    def make_inference(self, request):

        # ML MODEL

        predictions = self.ml_model.predict_proba(

            [

        x_data = np.array([

                [

                    request.n_packets_server_seconds,

                    request.n_packets_client_seconds,

                    request.n_packets_server_n_packets_client,

                    request.n_bits_server_n_bits_client,

                ]

            ]

        )

            ])

        predictions = self.model.run([self.label_name], {self.input_name: x_data.astype(np.float32)})[0]

        # Output format

        output_message = {

            "confidence": None,

            "time_start": request.time_start,

            "time_end": request.time_end,

        }

        if predictions[0][1] >= 0.5:

        if predictions[0] >= 0.5:

            output_message["confidence"] = predictions[0][1]

            output_message["tag_name"] = "Crypto"

            output_message["tag"] = 1