export onnx

import os, math, random, numpy
import shutil, json, argparse
import h5py, faiss
from tqdm import trange, tqdm
from PIL import Image

import torch
import torch.nn.functional as F
import torchvision
from torchvision import transforms
import onnx
import onnxruntime


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='NetVlad')
    parser.add_argument('--ckpt', type=str, default='best', help='resume from checkpoint.')
    parser.add_argument('--onnx', type=str, default='model.onnx', help='export to onnx.')
    options = parser.parse_args()

    model = torch.hub.load('yxgeee/OpenIBL', 'vgg16_netvlad', pretrained=True).eval()

    img = Image.open('image.jpg').convert('RGB') # modify the image path according to your need
    transformer = transforms.Compose([transforms.Resize((480, 640)), # (height, width)
                                    transforms.ToTensor(),
                                    transforms.Normalize(mean=[0.48501960784313836, 0.4579568627450961, 0.4076039215686255],
                                                        std=[0.00392156862745098, 0.00392156862745098, 0.00392156862745098])])
    img = transformer(img)
    output = model(img)

    # Export the model
    torch.onnx.export(model, input, options.onnx, export_params = True, verbose = True,
                      # store the trained parameter weights inside the model file
                      opset_version=16,  # the ONNX version to export the model to
                      #do_constant_folding=True,  # whether to execute constant folding for optimization
                      input_names=['image:0'],  # the model's input names
                      output_names=['descriptor:0'],  # the model's output names
                      dynamic_axes={'image:0': {2: 'image_height', 3: "image_width"}}
                      )

---

Reference

[0] Self-supervising Fine-grained Region Similarities for Large-scale Image Localization
[1] https://zhuanlan.zhihu.com/p/169596514 [2] https://github.com/yxgeee/OpenIBL