coco
21 lines
604 B
Python
21 lines
604 B
Python
import clip
|
|
from PIL import Image
|
|
from torch import Tensor
|
|
import torch
|
|
|
|
# Export ImageEncoder of the CLIP to onnx model
|
|
if __name__ == '__main__':
|
|
device = "cpu"
|
|
# print(clip.available_models())
|
|
model, preprocess = clip.load("ViT-B/32", device=device, jit=False)
|
|
i = Image.open("../../image.jpg")
|
|
input_tensor: Tensor = preprocess(i).unsqueeze(0).to(device)
|
|
vit = model.visual
|
|
vit.eval()
|
|
|
|
onnx_filename = 'clip-image-encoder.onnx'
|
|
torch.onnx.export(vit, input_tensor, onnx_filename)
|
|
# python -m onnxsim clip-image-encoder.onnx clip-image-encoder-optimized.onnx
|
|
|
|
|