Files
insightface/deploy/model_slim.py
2020-11-06 13:59:21 +08:00

33 lines
900 B
Python

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import sys
import os
import argparse
import numpy as np
import mxnet as mx
parser = argparse.ArgumentParser(description='face model slim')
# general
parser.add_argument('--model',
default='../models/model-r34-amf/model,60',
help='path to load model.')
args = parser.parse_args()
_vec = args.model.split(',')
assert len(_vec) == 2
prefix = _vec[0]
epoch = int(_vec[1])
print('loading', prefix, epoch)
sym, arg_params, aux_params = mx.model.load_checkpoint(prefix, epoch)
all_layers = sym.get_internals()
sym = all_layers['fc1_output']
dellist = []
for k, v in arg_params.iteritems():
if k.startswith('fc7'):
dellist.append(k)
for d in dellist:
del arg_params[d]
mx.model.save_checkpoint(prefix + "s", 0, sym, arg_params, aux_params)