WongKinYiu / yolor

implementation of paper - You Only Learn One Representation: Unified Network for Multiple Tasks (https://arxiv.org/abs/2105.04206)
GNU General Public License v3.0
1.98k stars 524 forks source link

Unable to get the same output as detect.py when using torchscript #233

Closed Clara-liu closed 2 years ago

Clara-liu commented 2 years ago

Hi,

Thank you for a wonderful paper and models.

I'd like to export the model into torchscript for doing inference on torchserve. Below is my code for converting to torchscript:

import torch
import argparse
from models.models import Darknet

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, help='weights path')
    parser.add_argument('--config', type=str, help='config file')
    parser.add_argument('--out-file', type=str, help='path to output torchscript file')
    parser.add_argument('--device', type=str, default='gpu', help='config file')
    parser.add_argument('--img-size', type=int, default=1280, help='image size n*n')
    parser.add_argument('--batch-size', type=int, default=1, help='batch size')
    opt = parser.parse_args()
    device = torch.device('cuda:0') if opt.device == 'gpu' else torch.device('cpu')  # set device
    img_size = [opt.img_size, opt.img_size]
    model = Darknet(opt.config, opt.img_size)  # Load PyTorch model
    model.load_state_dict(torch.load(opt.weights, map_location=device)['model'])
    if device == 'gpu':
        model.half()
    else:
        model.float()
    model.to(device).eval()
    img = torch.zeros((opt.batch_size, 3, *img_size)).to(device)
    y = model(img, augment=False)  # dry run
    with torch.no_grad():
        ts = torch.jit.trace(model, img)
    ts.save(opt.out_file)

I am able to export the model. However, the model output is completelty different from the result from detect.py.

Below is how I used the exported model for inferece:

import torch
from utils.general import non_max_suppression

model = torch.jit.load('mytorchscript.pth')
pred = model(img.unsqueeze(0))[0]
ouput = non_max_suppression(pred)

Are there any obvious mistakes I've made when exporting the model? 因为对这方面还不是十分熟悉 提前多谢大佬!

Clara-liu commented 2 years ago

I figured out why, I did not normalise my pixel values to between 0-1. I'm closing the issue.