一、前言
最近有空,把之前的项目梳理记录一下,惠已惠人。我的demo代码
二、详情
人脸模型是在 pytorch 下训练的,工程文件用的是这个:MobileFaceNet_Tutorial_Pytorch
训练完成之后,先转为onnx模型并做简化,代码如下:
def export_onnx():
import onnx
parser = argparse.ArgumentParser()
#parser.add_argument('--weights', type=str, default=r'F:\demo\Pytorch_demo\yolov5_fruit\weights\fruit_last.pt', help='weights path')
parser.add_argument('--img-size', nargs='+', type=int, default=[112, 112], help='image size')
parser.add_argument('--batch-size', type=int, default=1, help='batch size')
opt = parser.parse_args()
print(opt)
# Parameters
f = "zeng_mobileface_model\zeng_insightface.onnx"
sim_onnx_path = "zeng_mobileface_model\zeng_insightface_sim.onnx"
img = torch.zeros((opt.batch_size, 3, *opt.img_size)) # image size, (1, 3, 320, 192) iDetection
# Load pytorch model
#google_utils.attempt_download(opt.weights)
#model = torch.load(opt.weights, map_location=torch.device('cpu'))['model'].float()
#model.eval()
device = "cpu"
model = MobileFaceNet(512).to(device) # embeding size is 512 (feature vector)
model.load_state_dict(
torch.load('F:\demo\Pytorch_demo\MobileFaceNet_Tutorial_Pytorch\zeng_mobileface_model\Iter_455000_model.ckpt',
map_location='cpu')['net_state_dict'])
print('MobileFaceNet face detection model generated')
model.eval()
_ = model(img) # dry run
torch.onnx.export(model, img, f, verbose=False, opset_version=11, input_names=['images'],
output_names=['output']) # output_names=['classes', 'boxes']
# Check onnx model
model = onnx.load(f) # load onnx model
onnx.checker.check_model(model) # check onnx model
print(onnx.helper.printable_graph(model.graph)