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?
因为对这方面还不是十分熟悉 提前多谢大佬!
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:
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:
Are there any obvious mistakes I've made when exporting the model? 因为对这方面还不是十分熟悉 提前多谢大佬!