#!/usr/bin/env python
|
# -*- coding: utf-8 -*-
|
# @Time : 2020/10/26 15:50
|
# @Author : Scheaven
|
# @File : predictor.py
|
# @description:
|
import cv2, io
|
import torch
|
import torch.nn.functional as F
|
|
from modeling.meta_arch.build import build_model
|
from utils.checkpoint import Checkpointer
|
|
import onnx
|
import torch
|
from onnxsim import simplify
|
from torch.onnx import OperatorExportTypes
|
|
class ReID_Model(object):
|
def __init__(self, cfg):
|
self.cfg = cfg
|
|
self.predictor = DefaultPredictor(cfg)
|
|
def run_on_image(self, original_image):
|
|
predictions = self.predictor(original_image)
|
return predictions
|
|
def torch2onnx(self):
|
predictions = self.predictor.to_onnx()
|
|
class DefaultPredictor:
|
|
def __init__(self, cfg):
|
self.cfg = cfg.clone()
|
self.cfg.defrost()
|
self.cfg.MODEL.BACKBONE.PRETRAIN = False
|
self.model = build_model(self.cfg)
|
for param in self.model.parameters():
|
param.requires_grad = False
|
self.model.cuda()
|
self.model.eval()
|
|
Checkpointer(self.model).load(cfg.MODEL.WEIGHTS)
|
|
def __call__(self, image):
|
with torch.no_grad(): # https://github.com/sphinx-doc/sphinx/issues/4258
|
images = image.cuda()
|
self.model.eval()
|
predictions = self.model(images)
|
torch.set_printoptions(edgeitems=2048)
|
print("------------\n", predictions)
|
pred_feat = F.normalize(predictions)
|
pred_feat = pred_feat.cpu().data
|
return pred_feat
|
|
def to_onnx(self):
|
inputs = torch.randn(1, 3, self.cfg.INPUT.SIZE_TEST[0], self.cfg.INPUT.SIZE_TEST[1]).cuda()
|
onnx_model = self.export_onnx_model(self.model, inputs)
|
|
model_simp, check = simplify(onnx_model)
|
|
model_simp = self.remove_initializer_from_input(model_simp)
|
|
assert check, "Simplified ONNX model could not be validated"
|
|
onnx.save_model(model_simp, f"fastreid.onnx")
|
|
|
def export_onnx_model(self, model, inputs):
|
assert isinstance(model, torch.nn.Module)
|
def _check_eval(module):
|
assert not module.training
|
|
model.apply(_check_eval)
|
|
with torch.no_grad():
|
with io.BytesIO() as f:
|
torch.onnx.export(
|
model,
|
inputs,
|
f,
|
operator_export_type=OperatorExportTypes.ONNX_ATEN_FALLBACK,
|
)
|
onnx_model = onnx.load_from_string(f.getvalue())
|
|
# Apply ONNX's Optimization
|
all_passes = onnx.optimizer.get_available_passes()
|
passes = ["extract_constant_to_initializer", "eliminate_unused_initializer", "fuse_bn_into_conv"]
|
assert all(p in all_passes for p in passes)
|
onnx_model = onnx.optimizer.optimize(onnx_model, passes)
|
return onnx_model
|
|
|
def remove_initializer_from_input(self, model):
|
if model.ir_version < 4:
|
print(
|
'Model with ir_version below 4 requires to include initilizer in graph input'
|
)
|
return
|
|
inputs = model.graph.input
|
name_to_input = {}
|
for input in inputs:
|
name_to_input[input.name] = input
|
|
for initializer in model.graph.initializer:
|
if initializer.name in name_to_input:
|
inputs.remove(name_to_input[initializer.name])
|
|
return model
|
|
|
