After converting to TRT, the accuracy dropped significantly compared to PyTorch for the occlusion task
All the conversion steps are as follows.
step 1:
First, we modify the forward section in the following path: lib/models/litetrack/litetrack.py
def forward(self,template_feats, search ):
x = self.backbone(z=template_feats,x=search, mode='x')
out = self.forward_head(x)
return out['score_map'], out['size_map'], out['offset_map']
step2:
In the second step, we run the code for converting to ONNX.
import os
import sys
import torch
import torch.nn as nn
import importlib
from lib.train.data.processing_utils import sample_target
from lib.test.tracker.data_utils import Preprocessor
from lib.utils.box_ops import box_xywh_to_xyxy
from lib.train.data.processing_utils import transform_image_to_crop
import numpy as np
prj_path = os.path.join(os.path.dirname(__file__), '..')
if prj_path not in sys.path:
sys.path.append(prj_path)
class LiteTrackONNXWrapper(nn.Module):
def __init__(self, model,template_tensor,template_bbox,search_tensor):
super().__init__()
self.model = model
def forward(self, template_tensor, template_bbox, search_tensor):
device = next(self.model.parameters()).device # Ensure everything runs on the same device
template_feat = self.model.forward_z(template_tensor, template_bb=template_bbox)
out_dict = self.model(template_feats=template_feat, search=search_tensor)
print(len(out_dict))
return out_dict
def export_litetrack_full_onnx():
device = "cpu"
yaml_path = 'C:/Users/DEZH/Desktop/Litetrack_onnx_new/experiments/litetrack/B4_cae_center_all_ep300.yaml'
checkpoint_path = 'C:/Users/DEZH/Desktop/Litetrack_onnx_new/snapshot/LiteTrack_ep0300.pth.tar'
config_module = importlib.import_module('lib.config.litetrack.config')
config_module.update_config_from_file(yaml_path)
cfg = config_module.cfg
model_module = importlib.import_module('lib.models.litetrack.litetrack')
model = model_module.build_LiteTrack(cfg, training=False)
model.load_state_dict(torch.load(checkpoint_path, map_location='cpu')['net'], strict=True)
model.eval().to(device)
bs = 1
z_sz = cfg.DATA.TEMPLATE.SIZE
x_sz = cfg.DATA.SEARCH.SIZE
search_size = cfg.DATA.SEARCH.SIZE
template = torch.randn(bs, 3, z_sz, z_sz).to(device)
template_bb = torch.tensor([[0.5, 0.5, 0.5, 0.5]]).to(device)
search = torch.randn(bs, 3, x_sz, x_sz).to(device)
wrapper = LiteTrackONNXWrapper(model,template,template_bb,search).to(device)
wrapper.eval()
torch.onnx.export(
wrapper,
(template, template_bb, search),
"litetrack_new.onnx",
input_names=["template", "template_bb", "search"],
output_names=["out_dict"],
opset_version=11,
export_params=True,
)
if __name__ == "__main__":
export_litetrack_full_onnx()
step 3:
We convert the ONNX file to TRT (TensorRT) with this CLI
trtexec --onnx=./litetrack_new.onnx --saveEngine=./litetrack.trt --buildOnly --fp16
step 4:
After obtaining the file, we perform the build step.
we use this code main.cpp , litetrack.cpp litetrack.hpp
main.txt
litetrack.txt
Do you have any idea why the accuracy has dropped, especially since it is much lower compared to the accuracy of OSTrack TRT?
After converting to TRT, the accuracy dropped significantly compared to PyTorch for the occlusion task
All the conversion steps are as follows.
step 1:
First, we modify the forward section in the following path: lib/models/litetrack/litetrack.py
step2:
In the second step, we run the code for converting to ONNX.
step 3:
We convert the ONNX file to TRT (TensorRT) with this CLI
trtexec --onnx=./litetrack_new.onnx --saveEngine=./litetrack.trt --buildOnly --fp16
step 4:
After obtaining the file, we perform the build step.
we use this code main.cpp , litetrack.cpp litetrack.hpp
main.txt
litetrack.txt
Do you have any idea why the accuracy has dropped, especially since it is much lower compared to the accuracy of OSTrack TRT?