def __init__(self, checkpoint_path, device,
img_mean=np.array([128, 128, 128], dtype=np.float32),
img_scale=np.float32(1/255),
use_tensorrt=False):
from models.with_mobilenet import PoseEstimationWithMobileNet
from modules.load_state import load_state
self.img_mean = img_mean
self.img_scale = img_scale
self.device = 'cpu'
if device != 'CPU':
if torch.cuda.is_available():
self.device = torch.device('cuda:0')
else:
print('No CUDA device found, inferring on CPU')
net = PoseEstimationWithMobileNet()
checkpoint = torch.load(checkpoint_path, map_location='cpu')
if use_tensorrt:
from torch2trt import TRTModule
net = TRTModule()
net.load_state_dict(checkpoint)
else:
load_state(net, checkpoint)
net = net.to(self.device)
net.eval()
self.net = net
# Optimizes a model
# Relies on the platform => Must be compiled for every device
# The first command lin parameter must be the path to the old model
import sys
import torch
# import torch2trt
from models.with_mobilenet import PoseEstimationWithMobileNet
from torch2trt.torch2trt import *
MODEL_NAME = sys.argv[1]
NEW_MODEL_NAME = MODEL_NAME.replace(".pth", "_opt.pth")
WIDTH = 224
HEIGHT = 224
net = PoseEstimationWithMobileNet()
net.load_state_dict(torch.load(sys.argv[1]), strict=False)
net.eval().cuda()
data = torch.ones((1, 3, HEIGHT, WIDTH)).cuda()
model_trt = torch2trt(net, [data], fp16_mode=True, max_workspace_size=1 << 25)
torch.save(model_trt.state_dict(), NEW_MODEL_NAME)