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 torch2trt

device = '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_keys
    unused_pretrained_keys = ckpt_keys - model_keys
    missing_keys = model_keys - ckpt_keys
    print('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 True

def 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 x
    return {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 RuntimeError

    os.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 torch2trt


current_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_keys
    unused_pretrained_keys = ckpt_keys - model_keys
    missing_keys = model_keys - ckpt_keys
    print('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 True

def 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 x
    return {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 torch2trt

data = 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_face18

device = '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