tufei/torchscript-ssd

Export/tracing to libtorch

Closed this issue · 0 comments

Hi, I'm trying to convert a PyTorch model to TorchScript from your other repo pythorch-ssd, but in c++ with libtorch I can't detect anything, below my tracing script:

from vision.ssd.vgg_ssd import create_vgg_ssd, create_vgg_ssd_predictor
from vision.ssd.mobilenetv1_ssd import create_mobilenetv1_ssd, create_mobilenetv1_ssd_predictor
import cv2
import sys
import torch
import torchvision

example = torch.randn(1, 3, 300, 300).cuda()
label_path = "voc_labels.txt"
class_names = [name.strip() for name in open(label_path).readlines()]
#model = create_mobilenetv1_ssd(len(class_names), is_test=True).cuda()
model = create_vgg_ssd(len(class_names), is_test=True).cuda()
model.load('models/vgg16-ssd-mp-0_7726.pth')
model.eval().cuda()

# Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing.
traced_script_module = torch.jit.trace(model, example)
output = traced_script_module(example)
print(traced_script_module)
traced_script_module.save("traced_model_ssd_voc.pt")

EDIT: Ok I found your tracing code inside run_ssd_example.py