Skip to main content

Inference and Evaluation

This page covers how to bring your existing PyTorch or ONNX models onto the Bitfount platform specifically for inference or evaluation tasks. These are tasks that specifically use the bitfount.ModelInference or bitfount.ModelEvaluation algorithms. Models that do not perform any training or fine-tuning can be substantially simpler to onboard to Bitfount.

Bitfount supports two types of models:

  • PyTorch models
  • ONNX models

Regardless of the type of model you are bringing onto the Bitfount platform, the only thing you need to do is to make sure that your model implements the InferrableModelProtocol interface for inference or the EvaluableModelProtocol interface for evaluation. Both protocols require the implementation of initialise_model, which is used to initialise the model and deserialize, which is used to save the model parameters to a file. You will then need to additionally implement either the predict method for inference or the evaluate method for evaluation. Or both if you want to support both inference and evaluation.

tip

To validate that your model has implemented the appropriate interface correctly, you can use an isinstance check:

assert isinstance(MyModel(), InferrableModelProtocol)

PyTorch​

To make migration easier we've provided two base classes for PyTorch inference models:

  • PytorchInferenceModel
  • PytorchLightningInferenceModel

These base classes do not need to be used but may make migration easier as they provide a lot of the functionality you need to get started. For examples of how to implement a PyTorch inference model using these base classes, see the PyTorch inference model tutorial.

ONNX​

Open Neural Network eXchange (ONNX) is a popular open standard format for representing machine learning models across many popular frameworks such as PyTorch, TensorFlow, and scikit-learn. Models are typically not written in pure ONNX but rather in a framework-specific language such as PyTorch or TensorFlow and then converted to ONNX. This is what we also recommend you do in order to bring your model to Bitfount in ONNX format.

Converting your model to ONNX​

For more information on how to convert your model to ONNX, see the relevant documentation for your framework or the ONNX documentation. PyTorch for instance has a built-in support for converting models to ONNX [documentation] whereas TensorFlow requires the use of a dedicated library for converting models to ONNX [documentation].

A simple example of converting a PyTorch model to ONNX is shown below:

convert_to_onnx.py
import torchfrom torch import nn as nnclass BinaryClassificationModel(nn.Module):    """Simple binary classification model with 2 input features."""    def __init__(self, in_features: int = 2) -> None:        """Initialise the binary classification model.        Args:            in_features: Number of input features (default: 2 for A and B).        """        super().__init__()        self.linear = nn.Linear(in_features, 1)        self.sigmoid = nn.Sigmoid()    def forward(self, x: torch.Tensor) -> torch.Tensor:        """Forward pass through the model.        Args:            x: Input tensor of shape (batch_size, in_features).        Returns:            Output tensor of shape (batch_size, 1) with sigmoid activation.        """        x = self.linear(x)        x = self.sigmoid(x)        return xif __name__ == "__main__":    model = BinaryClassificationModel()    model.eval()    # Create dummy input for tracing    dummy_input = torch.randn(1, 2)  # 2 features    # Export to ONNX    model_path = "binary_classification.onnx"    torch.onnx.export(        model,        (dummy_input,),        str(model_path),        input_names=["x"],        output_names=["y"],        export_params=True,        opset_version=13,        dynamic_axes={"x": {0: "batch_size"}, "y": {0: "batch_size"}},    )

Encapsulating your model​

Things work a little differently for ONNX models due to the fact that ONNX serialises the entire model graph alongside the weights. This means that you don't actually need to display your underlying model code when you upload your model to Bitfount. However, you still need to encapsulate your model with the appropriate InferrableModelProtocol or EvaluableModelProtocol interface so that Bitfount is able to interact with it.

tip

Bitfount allows models to be uploaded either publicly or privately. However, even a private model will display the underlying model code when access has been granted to collaborators. If you want to keep your model code private, we recommend converting your model to an ONNX model.

To make things easier, we have provided a base class for ONNX inference models: ONNXModel. This base class implements much of the boilerplate required to get started. See the API documentation for more details on the ONNXModel class.

Using this class could make your ONNX model code look as simple as:

binary_classification_onnx_model.py
from bitfount.backends.onnx.models import ONNXModelclass BinaryClassificationModel(ONNXModel):    """Binary classification model using ONNX."""

The code is not missing - for most use cases there is simply no need to implement or override anything unless the input data or the output values require special handling. The ONNXModel class implements all of the required functionality. And anything specific to your model architecture has been serialised alongside the model weights.

For more information on how to upload the model code and .onnx file to Bitfount, see the Uploading your model page. The .onnx file is the file that contains the model graph alongside the model weights.