Jetson学习笔记(四):pth(torch模型文件)转trt(tensorrt引擎文件)实操

install torch2trt

git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt
sudo python setup.py install --plugins

具体代码1

from retinaface.models.retinaface import RetinaFace, PriorBox  # 导入网络
import torch,os
from torch2trt import torch2trtdevice = 'cuda' if torch.cuda.is_available() else 'cpu'
current_dir=os.path.dirname(os.path.abspath(__file__)) # 获取当前路径cfg = {'name': 'mobilenet0.25','min_sizes': [[16, 32], [64, 128], [256, 512]],'steps': [8, 16, 32],'variance': [0.1, 0.2],'clip': False,'loc_weight': 2.0,'gpu_train': True,'batch_size': 32,'ngpu': 1,'epoch': 250,'decay1': 190,'decay2': 220,'image_size': 640,'pretrain': True,'return_layers': {'stage1': 1, 'stage2': 2, 'stage3': 3},'in_channel': 32,'out_channel': 64
}def load_model(model, pretrained_path, device):print('Loading pretrained model from {}'.format(pretrained_path))pretrained_dict = torch.load(pretrained_path, map_location=device)if "state_dict" in pretrained_dict.keys():pretrained_dict = remove_prefix(pretrained_dict['state_dict'], 'module.')else:pretrained_dict = remove_prefix(pretrained_dict, 'module.')check_keys(model, pretrained_dict)model.load_state_dict(pretrained_dict, strict=False)def check_keys(model, pretrained_state_dict):ckpt_keys = set(pretrained_state_dict.keys())model_keys = set(model.state_dict().keys())used_pretrained_keys = model_keys & ckpt_keysunused_pretrained_keys = ckpt_keys - model_keysmissing_keys = model_keys - ckpt_keysprint('Missing keys:{}'.format(len(missing_keys)))print('Unused checkpoint keys:{}'.format(len(unused_pretrained_keys)))print('Used keys:{}'.format(len(used_pretrained_keys)))assert len(used_pretrained_keys) > 0, 'load NONE from pretrained checkpoint'return Truedef remove_prefix(state_dict, prefix):''' Old style model is stored with all names of parameters sharing common prefix 'module.' '''print('remove prefix \'{}\''.format(prefix))f = lambda x: x.split(prefix, 1)[-1] if x.startswith(prefix) else xreturn {f(key): value for key, value in state_dict.items()}def create_engine(weights, device, eps=1e-3):print("Create trt engine for retintaface...")model = RetinaFace(cfg).to(device)load_model(model, weights, device)model.eval()x = torch.ones((1, 3, cfg["image_size"], cfg["image_size"])).to(device)  # cfg["image_size"=640 根据自己的模型输出设置的大小model_trt = torch2trt(model, [x])print("Ok. Check outputs...")y = model(x)y_trt = model_trt(x)for out, out_trt in zip(y, y_trt):if torch.max(torch.abs(out - out_trt)) > eps:raise RuntimeErroros.makedirs(os.path.join(current_dir, "engines"), exist_ok=True)torch.save(model_trt.state_dict(), os.path.join(current_dir, "engines", f"retina_trt_{device}.trt"))print('trt create finish.......')return model_trt

运行结果

具体代码2

from arcface.resnet import resnet_face18
import torch,os
from torch2trt import torch2trtcurrent_dir=os.path.dirname(os.path.abspath(__file__)) # 获取当前路径def load_model(model, pretrained_path, device):print('Loading pretrained model from {}'.format(pretrained_path))pretrained_dict = torch.load(pretrained_path, map_location=device)if "state_dict" in pretrained_dict.keys():pretrained_dict = remove_prefix(pretrained_dict['state_dict'], 'module.')else:pretrained_dict = remove_prefix(pretrained_dict, 'module.')check_keys(model, pretrained_dict)model.load_state_dict(pretrained_dict)def check_keys(model, pretrained_state_dict):ckpt_keys = set(pretrained_state_dict.keys())model_keys = set(model.state_dict().keys())used_pretrained_keys = model_keys & ckpt_keysunused_pretrained_keys = ckpt_keys - model_keysmissing_keys = model_keys - ckpt_keysprint('Missing keys:{}'.format(len(missing_keys)))print('Unused checkpoint keys:{}'.format(len(unused_pretrained_keys)))print('Used keys:{}'.format(len(used_pretrained_keys)))assert len(used_pretrained_keys) > 0, 'load NONE from pretrained checkpoint'return Truedef remove_prefix(state_dict, prefix):''' Old style model is stored with all names of parameters sharing common prefix 'module.' '''print('remove prefix \'{}\''.format(prefix))f = lambda x: x.split(prefix, 1)[-1] if x.startswith(prefix) else xreturn {f(key): value for key, value in state_dict.items()}def create_engine(weights, device, eps=1e-3):print("Create trt engine for retintaface...")model = resnet_face18(use_se=True).cuda()load_model(model, weights, device)model.eval()x = torch.ones((1, 1,128,128)).cuda()  # cfg["image_size"=640 根据自己的模型输出设置的大小model_trt = torch2trt(model, [x])print('save')os.makedirs(os.path.join(current_dir, "engines"), exist_ok=True)torch.save(model_trt.state_dict(), os.path.join(current_dir, "engines", f"arcface_trt_256.trt"))print('trt create finish.......')return model_trt
device = 'cuda' if torch.cuda.is_available() else 'cpu'
weights='/home/lqs/Documents/Engineering_CYB/pth_onnx_model/resnet18_256_90.pth'
create_engine(weights, device, eps=1e-3)

运行结果

Create trt engine for retintaface...
Loading pretrained model from /home/lqs/Documents/Engineering_CYB/pth_onnx_model/resnet18_256_90.pth
remove prefix 'module.'
Missing keys:0
Unused checkpoint keys:0
Used keys:221
save
trt create finish.......

具体代码3

# -*- coding: utf-8 -*-
import torchvision
import torch
from torch2trt import torch2trtdata = torch.randn((1, 3, 224, 224)).cuda().half()
model = torchvision.models.resnet18(pretrained=True).cuda().half().eval()
output = model(data)# pytorch -> tensorrt
model_trt = torch2trt(model, [data], fp16_mode=True)
output_trt = model_trt(data)# compare
print('max error: %f' % float(torch.max(torch.abs(output - output_trt))))
print("mse :%f" % float((output - output_trt)**2))# save tensorrt model
torch.save(model_trt.state_dict(), "resnet18_trt.pth")# load tensorrt model
from torch2trt import TRTModule
model_trt = TRTModule()
model_trt.load_state_dict(torch.load('resnet18_trt.pth'))

# -*- coding: utf-8 -*-
import torchvision
import torch
from collections import OrderedDict
from torch2trt import torch2trt
from arcface.resnet import resnet_face18device = 'cuda' if torch.cuda.is_available() else 'cpu'
data = torch.randn((1, 1, 128, 128)).cuda()
model = resnet_face18(use_se=True).cuda()
model_path = '/home/lqs/Documents/Engineering_CYB/pth_onnx_model/resnet18_256_90.pth'
state_dict = torch.load(model_path, map_location=device)
print(1)
mew_state_dict = OrderedDict()
model_dict = model.state_dict()
pretrained_dict = {k: v for k, v in state_dict.items() if (k in model_dict and 'fc' not in k)}
model_dict.update(pretrained_dict)
print(2)
model.load_state_dict(model_dict)
model.eval()
print(3)
output = model(data)
print(4)
# pytorch -> tensorrt
model_trt = torch2trt(model, [data], fp16_mode=True)
print('begin to save')
# save tensorrt model
torch.save(model_trt.state_dict(), "arcface_trt_256.trt")

运行结果

1
2
3
4
begin to save