DAB-DETR icon indicating copy to clipboard operation
DAB-DETR copied to clipboard

Convert DAB-deformable-DETR to ONNX

Open sazani opened this issue 2 years ago • 1 comments

I am trying to convert the generated model that I trained and also your pretrained model to ONNX but unfortunately I faced the following error message:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0, and cpu! (when checking argument for argument index in method wrapper__index_select) By the way I have used static and dynamic input and I have used the following code:

import torch.onnx import os, sys import torch import numpy as np

from models import build_dab_deformable_detr from util.slconfig import SLConfig import torchvision import torchvision.transforms.functional as TF

from PIL import Image import transforms as T

import cv2 import argparse device = torch.device('cuda:0' )

if name == "main": parser = argparse.ArgumentParser() parser.add_argument('--model_checkpoint_path', help="change the path of the model checkpoint.", default="./Checkpoints/checkpoint.pt") parser.add_argument('--model_config_path', help="change the path of the model config file", default="./Checkpoints/config.json") args = parser.parse_args() model_config_path = args.model_config_path model_checkpoint_path = args.model_checkpoint_path args_config = SLConfig.fromfile(model_config_path) model, criterion, post_processors = build_dab_deformable_detr(args_config) checkpoint = torch.load(model_checkpoint_path, map_location=device) model.load_state_dict(checkpoint['model']) model = model.to(device) img_size =[1080,1920] input = torch.zeros(1, 3, *img_size) input = input.to(device) model.eval() results =model(input) torch.onnx.export( model, input, "test.onnx", input_names=["input"], output_names=["output"], export_params=True, opset_version=11, # I have also tried version 12,13,14,15 # dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}, # shape(1,3,640,640) # 'output': {0: 'batch', 1: 'anchors'} # shape(1,25200,85) # } ,#if dynamic else None dynamic_axes = None, )

sazani avatar Dec 27 '22 03:12 sazani