Shiaoming / ALIKED

ALIKED: A Lighter Keypoint and Descriptor Extraction Network via Deformable Transformation
https://arxiv.org/pdf/2304.03608.pdf
BSD 3-Clause "New" or "Revised" License
190 stars 16 forks source link

How to convert the model to torch script file ? #4

Closed Pepper-FlavoredChewingGum closed 1 year ago

Pepper-FlavoredChewingGum commented 1 year ago

Thank you for your excellent works. @Shiaoming

I want to convert the model (aliked-n16rot.pth) to a torch script file so that the ALIKED can be used directly in C + +. But I ran into a lot of issues during the export process. Here's my code for a pt file; could you help me?

Thank you very much for your assistance. I eagerly await your response.

`import torch, cv2 import argparse import logging from nets.aliked import ALIKED from torchvision.transforms import ToTensor

def model_with_run(model): def fn(image): results = model.run(image) return results return fn

parser = argparse.ArgumentParser(description='ALIKED image pair Demo.') parser.add_argument('--model', choices=['aliked-t16', 'aliked-n16', 'aliked-n16rot', 'aliked-n32'], default="aliked-n16rot", help="The model configuration") parser.add_argument('--device', type=str, default='cuda', help="Running device (default: cuda).") parser.add_argument('--top_k', type=int, default=-1, help='Detect top K keypoints. -1 for threshold based mode, >0 for top K mode. (default: -1)') parser.add_argument('--scores_th', type=float, default=0.2, help='Detector score threshold (default: 0.2).') parser.add_argument('--n_limit', type=int, default=5000, help='Maximum number of keypoints to be detected (default: 5000).') args = parser.parse_args()

logging.basicConfig(level=logging.INFO)

model = ALIKED(model_name=args.model, device=args.device, top_k=args.top_k, scores_th=args.scores_th, n_limit=args.n_limit) model = model_with_run(model)

img_rgb = cv2.imread('1311868169.163498.png') img_rgb = cv2.cvtColor(img_rgb, cv2.COLOR_BGR2RGB) img_tensor = ToTensor()(img_rgb) img_tensor = imgtensor.to('cuda').unsqueeze(0)

traced_model = torch.jit.trace(model, img_tensor)

traced_model.save('ALIKED.pt') `

Shiaoming commented 1 year ago

you need to take care of the custom_ops and deformable conv

adricostas commented 1 year ago

Hello! I'm also interested in using the model directly in C++. I was trying to convert the model to a TorchScript file via torch.jit.script but this implies to change several things in the code. Have you been able to convert it using torch.jit.trace finally? Thanks!

Shiaoming commented 1 year ago

I have successfull exported to torch scirpt.

adricostas commented 1 year ago

Nice, thanks. And are you planning to share an ALIKED-cpp version as well as you did with ALIKE?

Shiaoming commented 1 year ago

Nope