mahan-ym/bioclip-2-quantized

Model Card for BioCLIP-2 Quantized

BioCLIP-2 Quantized is a quantized version of BioCLIP-2, a foundation model for biological organismal images read more about the model here. The quantized model is designed to reduce memory usage and improve inference with cost of losing a little accuracy.

Model Description

BioCLIP-2 Quantized is dynamically quantized from the original BioCLIP-2 model. The Pytorch onnx.export function is used to convert the original model to ONNX format, and then the onnxruntime.quantization.quantize_dynamic function is used to perform dynamic quantization on the ONNX model with weight_type set to QuantType.QInt8.

Quantization Details

You can find the quantization script here.

Inference

To use the quantized model for inference, you can use the onnxruntime library. Here is an example code snippet:

import onnxruntime as ort
import torch
import torch.nn.functional as F
import numpy as np
from huggingface_hub import hf_hub_download


# Load the quantized model
ort_session = ort.InferenceSession("path/to/bioclip-2-quantized.onnx", providers=['CPUExecutionProvider'])

# only return one label
k = 1

# Preprocess image
img_tensor = preprocess_img(img).unsqueeze(0)
img_np = img_tensor.numpy()

# Run ONNX inference
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

img_features_np = session.run([output_name], {input_name: img_np})[0]

# Convert back to torch for compatibility with existing code
img_features = torch.from_numpy(img_features_np)
img_features = F.normalize(img_features, dim=-1)

# optional: for the open ended classification that you need a text embedding:
txt_emb = torch.from_numpy(
    np.load(
        hf_hub_download(
            repo_id="imageomics/TreeOfLife-200M",
            filename="embeddings/txt_emb_species.npy",
            repo_type="dataset",
        )
    )
)
# or you can skip this and use zero-shot classification with your own text inputs

# Use the same text embeddings and logit scale from the original model (logit_scale of the main model: 100.00000762939453)
logits = (model.logit_scale.exp() * img_features @ txt_emb).squeeze()
probs = F.softmax(logits, dim=0)

topk = probs.topk(k)
prediction_dict = {
    format_name(*txt_names[i]): prob
    for i, prob in zip(topk.indices, topk.values)
}

print(prediction_dict)

Tradeoff

The Model got tested on Animals and Plants dataset by Nguyen Le Truong Thien for open-ended species classification and the highest probability class is selected as the predicted class and compared with the main BioCLIP-2 model. The results are as follows:

Acknowledgements

The main model is developed by the Imageomics Institute team. And the current model is just a quantized version of the main model to reduce memory usage and improve inference speed and make the model more accessible.