<jupyter_start><jupyter_text>[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/cene555/ruCLIP-SB/blob/main/notebooks/ruCLIP_SB_onnx.ipynb)<jupyter_code>#@title Allowed Resources
import multiprocessing
import torch
from psutil import virtual_memory

ram_gb = round(virtual_memory().total / 1024**3, 1)

print('CPU:', multiprocessing.cpu_count())
print('RAM GB:', ram_gb)
print("PyTorch version:", torch.__version__)
print("CUDA version:", torch.version.cuda)
print("cuDNN version:", torch.backends.cudnn.version())
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("device:", device.type)

!nvidia-smi<jupyter_output>CPU: 2
RAM GB: 12.7
PyTorch version: 1.10.0+cu111
CUDA version: 11.1
cuDNN version: 8005
device: cuda
Tue Jan 25 17:45:47 2022       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 495.46       Driver Version: 460.32.03    CUDA Version: 11.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla K80           Off  | 00000000:00:04.0 Off |                    0 |
| N/A   34C    P8    28W / 149W |      3MiB / 11441MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------[...]<jupyter_text>Install requirements<jupyter_code>%%capture
!pip install git+https://github.com/cene555/ruCLIP-SB.git
!pip install pymorphy2
!gdown -O ruCLIP-SB.pkl https://drive.google.com/uc?id=1-CghuC9TCIDyn5H3zQS6ho_TNiudzJCX

!pip install git+https://github.com/Lednik7/CLIP-ONNX.git
!pip install git+https://github.com/openai/CLIP.git
!pip install onnxruntime-gpu

!wget -c -O CLIP.png https://github.com/openai/CLIP/blob/main/CLIP.png?raw=true<jupyter_output><empty_output><jupyter_text>Import libraries<jupyter_code>import torch
from torchvision import transforms
import transformers
from transformers import BertTokenizer
from ruclipsb import ruCLIPSB
from ruclipsb.utils import tokenize, _convert_image_to_rgb
from PIL import ImageCms, Image
import cv2
import numpy as np
try:
    from torchvision.transforms import InterpolationMode
    BICUBIC = InterpolationMode.BICUBIC
except ImportError:
    BICUBIC = Image.BICUBIC
import warnings

warnings.filterwarnings("ignore", category=UserWarning)
torch.manual_seed(1)
device = torch.device('cpu')<jupyter_output><empty_output><jupyter_text>Load RuCLIP-SB model<jupyter_code>model = ruCLIPSB()
model.load_state_dict(torch.load('ruCLIP-SB.pkl', map_location=device))
model = model.half().to(device)

model = model.eval()
for x in model.parameters(): x.requires_grad = False
torch.cuda.empty_cache()
tokenizer = BertTokenizer.from_pretrained("cointegrated/rubert-tiny")
transform = transforms.Compose([
            transforms.Resize(224),
            transforms.CenterCrop(224),
            _convert_image_to_rgb,
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225]),])<jupyter_output><empty_output><jupyter_text>Prepare functions<jupyter_code># batch first
image = transform(Image.open("CLIP.png")).unsqueeze(0).cpu() # [1, 3, 224, 224]
image_onnx = image.detach().cpu().numpy().astype(np.float32)

# batch first
texts = ['диаграмма', 'собака', 'кошка']
text_tokens, attention_mask = tokenize(tokenizer, texts, 77)
text_tokens, attention_mask = text_tokens.cpu(), attention_mask.cpu() # [3, 77]
text_onnx = torch.stack([text_tokens, attention_mask]).detach().cpu().numpy().astype(np.int64)
class Textual(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model

    def forward(self, input_data):
        input_ids, attention_mask = input_data
        x = self.model.transformer(input_ids=input_ids, attention_mask=attention_mask)
        x = x.last_hidden_state[:, 0, :]
        x = self.model.final_ln(x)
        return x<jupyter_output><empty_output><jupyter_text>Convert RuCLIP-SB model to ONNX<jupyter_code>from clip_onnx import clip_onnx

def convert_textual(self, dummy_input):
    textual = Textual(self.model)
    torch.onnx.export(textual, dummy_input, self.textual_path,
                  input_names=['input'], output_names=['output'],
                  export_params=True, verbose=False, opset_version=14,
                  do_constant_folding=True,
                  dynamic_axes={'input': {1: 'batch_size'}, 'output': {0: 'batch_size'}})
    self.onnx_checker(self.textual_path)

clip_onnx.convert_textual = convert_textual

visual_path = "clip_visual.onnx"
textual_path = "clip_textual.onnx"

dummy_input_text = torch.stack([text_tokens, attention_mask]).detach().cpu()
onnx_model = clip_onnx(model.float().cpu(), visual_path=visual_path, textual_path=textual_path)
onnx_model.convert2onnx(image, dummy_input_text, verbose=True)<jupyter_output><empty_output><jupyter_text>[ONNX] CUDA inference mode<jupyter_code># ['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider']
onnx_model.start_sessions(providers=["CUDAExecutionProvider"]) # cuda mode
image_features = onnx_model.encode_image(image_onnx)
text_features = onnx_model.encode_text(text_onnx)

logits_per_image, logits_per_text = onnx_model(image_onnx, text_onnx)
probs = logits_per_image.softmax(dim=-1).detach().cpu().numpy()

print("Label probs:", probs) # [[0.9844646  0.01167088 0.00386453]]
%timeit onnx_model.encode_image(image_onnx)
%timeit onnx_model.encode_text(text_onnx)<jupyter_output>100 loops, best of 5: 2.76 ms per loop
{'fraction_non_alphanumeric': 0.17092436974789915, 'fraction_numerical': 0.03142857142857143, 'mean_word_length': 3.484551620195931, 'pattern_counts': {'":': 0, '<': 27, '<?xml version=': 0, '>': 27, 'https://': 7, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 1, 'word_list_counts': {'cursed_substrings.json': 7, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '661751', 'license': 'MIT', 'n_tokens_mistral': 2169, 'n_tokens_neox': 2011, 'n_words': 345, 'path': 'notebooks/ruCLIP_SB_onnx.ipynb', 'repo_name': 'cene555/ruCLIP-SB', 'size': '5950'}