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)