Patch clip model for ONNX compatibility

[chajath]

Team,

When productionizing AI models like CLIP, it is often useful to be able to export to ONNX so that we can utilize training and serving ecosystem built around ONNX runtime. This PR includes changes that are necessary to make ONNX compilation and runtime working without needing to further patch the code:

• Changes to use INT32 for tokenization, since ONNX doesn't yet support ArgMax(INT64). See [ONNXRuntimeError] : 9 : NOT_IMPLEMENTED : Could not find an implementation for the node ArgMax_1094:ArgMax(11) microsoft/onnxruntime#10068
• Use explicit dimension for norm

With these changes, we can now compile and run ONNX models. For example:

import clip
import torch

m, pre = clip.load("RN50")
npx = m.visual.input_resolution
dummy_image = torch.randn(10, 3, npx, npx)
dummy_texts = clip.tokenize(["quick brown fox", "lorem ipsum"])
m.forward(dummy_image ,dummy_texts) # Original CLIP result (1)

torch.onnx.export(m, (dummy_image, dummy_texts), "clip_resnet.onnx", export_params=True,
  input_names=["IMAGE", "TEXT"],
  output_names=["LOGITS_PER_IMAGE", "LOGITS_PER_TEXT"],
  opset_version=14,
  dynamic_axes={
      "IMAGE": {
          0: "image_batch_size",
      },
      "TEXT": {
          0: "text_batch_size",
      },
      "LOGITS_PER_IMAGE": {
          0: "image_batch_size",
          1: "text_batch_size",
      },
      "LOGITS_PER_TEXT": {
          0: "text_batch_size",
          1: "image_batch_size",
      },
  }
)

# Now run onnxruntime to verify
import onnxruntime as ort

ort_sess = ort.InferenceSession("clip_resnet.onnx")
result=ort_sess.run(["LOGITS_PER_IMAGE", "LOGITS_PER_TEXT"], 
  {"IMAGE": dummy_image.numpy(), "TEXT": dummy_texts.numpy()})
result # verify that result is comparable to (1)

I've locally verified the result. I see that CLIP has a github action to run pytest. I will be happy to further contribute by adding onnx-specific tests to see if the model can be compiled and the resulting model is correct.

I understand that ONNX compatibility might not be the primary goal of CLIP repo, but maintaining compatibility with ONNX will be immensely useful to the ML practitioners out in the wild, so please take a look at the change.

Thanks!

[chajath]

Regarding the test error, I see that torch 1.7.1 requires indices to be LongTensor, whereas 1.8 and above allows both LongTensor and IntTensor. Do you have suggestions as to how we should workaround? I wonder how long 1.7.1 is going to be supported by CLIP.

[chajath]

So I've pushed the change to preserve the behavior when torch version is < 1.8.0. This makes all the tests to pass again in 1.7.1, but 1.7.1 users won't benefit from this fix.

[hanxiao]

For ONNX support, please use https://github.com/jina-ai/clip-as-service/

[chajath]

For ONNX support, please use https://github.com/jina-ai/clip-as-service/

Respectfully, I don't think it's appropriate to promote your service in this PR.

[jongwook]

Thanks for the PR! I think branching based on torch version makes sense.

[deshwalmahesh]

Team,

When productionizing AI models like CLIP, it is often useful to be able to export to ONNX so that we can utilize training and serving ecosystem built around ONNX runtime. This PR includes changes that are necessary to make ONNX compilation and runtime working without needing to further patch the code:

• Changes to use INT32 for tokenization, since ONNX doesn't yet support ArgMax(INT64). See [ONNXRuntimeError] : 9 : NOT_IMPLEMENTED : Could not find an implementation for the node ArgMax_1094:ArgMax(11) microsoft/onnxruntime#10068
• Use explicit dimension for norm

With these changes, we can now compile and run ONNX models. For example:

import clip
import torch

m, pre = clip.load("RN50")
npx = m.visual.input_resolution
dummy_image = torch.randn(10, 3, npx, npx)
dummy_texts = clip.tokenize(["quick brown fox", "lorem ipsum"])
m.forward(dummy_image ,dummy_texts) # Original CLIP result (1)

torch.onnx.export(m, (dummy_image, dummy_texts), "clip_resnet.onnx", export_params=True,
  input_names=["IMAGE", "TEXT"],
  output_names=["LOGITS_PER_IMAGE", "LOGITS_PER_TEXT"],
  opset_version=14,
  dynamic_axes={
      "IMAGE": {
          0: "image_batch_size",
      },
      "TEXT": {
          0: "text_batch_size",
      },
      "LOGITS_PER_IMAGE": {
          0: "image_batch_size",
          1: "text_batch_size",
      },
      "LOGITS_PER_TEXT": {
          0: "text_batch_size",
          1: "image_batch_size",
      },
  }
)

# Now run onnxruntime to verify
import onnxruntime as ort

ort_sess = ort.InferenceSession("clip_resnet.onnx")
result=ort_sess.run(["LOGITS_PER_IMAGE", "LOGITS_PER_TEXT"], 
  {"IMAGE": dummy_image.numpy(), "TEXT": dummy_texts.numpy()})
result # verify that result is comparable to (1)

I've locally verified the result. I see that CLIP has a github action to run pytest. I will be happy to further contribute by adding onnx-specific tests to see if the model can be compiled and the resulting model is correct.

I understand that ONNX compatibility might not be the primary goal of CLIP repo, but maintaining compatibility with ONNX will be immensely useful to the ML practitioners out in the wild, so please take a look at the change.

Thanks!

How can you use model.encode_image() instead of model.forward() for the onnx model?