def __init__(self, model_path, long_side, network):
torch.set_grad_enabled(False)
if network == 'mobilenet':
self.cfg = cfg_mnet
net = RetinaFace(cfg=self.cfg, phase='test')
elif network == 'slim':
self.cfg = cfg_slim
net = Slim(cfg=self.cfg, phase='test')
elif network == 'RFB':
self.cfg = cfg_rfb
net = RFB(cfg=self.cfg, phase='test')
else:
print("not supported network!!")
exit(0)
self.net = load_model(net, model_path, True)
self.net.eval()
print("Finished loading model!")
cudnn.benchmark = True
self.device = torch.device("cpu")
self.net = self.net.to(self.device)
self.long_side = long_side
if __name__ == '__main__':
torch.set_grad_enabled(False)
cfg = None
net = None
if args.network == "mobile0.25":
cfg = cfg_mnet
net = RetinaFace(cfg = cfg, phase = 'test')
elif args.network == "slim":
cfg = cfg_slim
net = Slim(cfg = cfg, phase = 'test')
elif args.network == "RFB":
cfg = cfg_rfb
net = RFB(cfg = cfg, phase = 'test')
else:
print("Don't support network!")
exit(0)
net = load_model(net, args.trained_model, args.cpu)
net.eval()
print('Finished loading model!')
print(net)
cudnn.benchmark = True
device = torch.device("cpu" if args.cpu else "cuda")
net = net.to(device)
cap = cv2.VideoCapture(args.video_path)
# testing begin
args = parser.parse_args()
if not os.path.exists(args.save_folder):
os.mkdir(args.save_folder)
cfg = None
net = None
if args.network == "mobile0.25":
cfg = cfg_mnet
net = RetinaFace(cfg=cfg)
elif args.network == "slim":
cfg = cfg_slim
net = Slim(cfg=cfg)
elif args.network == "RFB":
cfg = cfg_rfb
net = RFB(cfg=cfg)
elif args.network == "efficientdet":
cfg = cfg_efficientdet
net = EfficientDet(cfg=cfg)
else:
print("Don't support network!")
exit(0)
print("Printing net...")
#print(net)
rgb_mean = (104, 117, 123) # bgr order
num_classes = 2
img_dim = cfg['image_size']
num_gpu = cfg['ngpu']
batch_size = cfg['batch_size']
#
# torch.set_grad_enabled(False)
# x_onnx=np.float32(x)
# for i in range(1000):
# t1_onnx = time.time()
# y_onnx = session.run([], {input_name: x_onnx})
# t1_onnxend = time.time()- t1_onnx
# onnx_time.append(t1_onnxend)
# average_time = np.sum(np.array(onnx_time[1:]))/999
# print("ONNX inference time:{}\n".format(average_time))
# print(y_onnx)
#Prepare models for torch and tensorRT
net = RFB(cfg=None, phase='test')
trained_model = "./weights/RBF_Final.pth"
device = torch.device("cuda")
original_torch_time = []
tensorRT_time = []
model = load_model(net, trained_model, True).eval().to(device)
x = torch.ones((1, 3, 224, 224)).cuda()
size = [(480, 600), (720, 1280), (1920, 1080)]
for sz in size:
print("***** SIZE {} *****".format(sz))
x = torch.rand((1, 3, sz[0], sz[1])).to(device)
x = x.to(device)
model_trt = torch2trt(model, [x])
