pytorch_classifier.py

from typing import List

import torch
import torch.nn as nn
import torch.nn.functional as F

from classifier import Classifier


PYTORCH_MODEL_INPUT_DIM = 4


class Model(nn.Module):
    def __init__(self, input_dim=PYTORCH_MODEL_INPUT_DIM):
        super(Model, self).__init__()
        self.layer1 = nn.Linear(input_dim, 50)
        self.layer2 = nn.Linear(50, 50)
        self.layer3 = nn.Linear(50, 3)

    def forward(self, x):
        x = F.relu(self.layer1(x))
        x = F.relu(self.layer2(x))
        x = F.softmax(self.layer3(x), dim=1)
        return x

class PytorchClassifier(Classifier):
    def __init__(self, model_path):
        super().__init__(model_path)
        self.model = Model()
        self.model.load_state_dict(torch.load(self.model_path))
        self.model.eval()

    def predict(self, input_data: List[List[float]]):
        with torch.no_grad():
            probas = self.model(torch.Tensor(input_data)).tolist()
        return self._prepare_response_dict(probas)