Open ChengLiu-xsy opened 6 months ago
import torch from lib.model_test import D2Net from lib.utils import preprocess_image from lib.pyramid import process_multiscale
use_cuda = torch.cuda.is_available() device = torch.device("cuda:0" if use_cuda else "cpu")
model_file = '/home/badal/d2-net/d2_tf.pth' # Specify the path to your model file use_relu = True # Specify whether to use ReLU
model = D2Net( model_file=model_file, use_relu=use_relu, use_cuda=use_cuda )
model.eval()
dummy_input = torch.randn(1, 3, 480, 640, device=device)
onnx_file_path = '/home/badal/d2-net/onnx/model_new.onnx'
torch.onnx.export(model, dummy_input, onnx_file_path, export_params=True, opset_version=11, input_names=['input'], output_names=['descriptor' , 'score' , 'keypoint' ])
print(f"Model has been converted to ONNX and saved to {onnx_file_path}")
hi: Could you please provide the source code of exporting onnx file?