Browse Source
Add distilled weights
Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
main
4 changed files with
4 additions and
4 deletions
-
nn_fingerprint.py
-
saved_model/nnfp_fma.pt
-
BIN
saved_model/nnfp_fma_distill_128.pt
-
BIN
saved_model/pfann_fma_s.pt
|
@ -60,7 +60,7 @@ class NNFingerprint(NNOperator): |
|
|
log.info('Loading model...') |
|
|
log.info('Loading model...') |
|
|
if model_path is None: |
|
|
if model_path is None: |
|
|
path = str(Path(__file__).parent) |
|
|
path = str(Path(__file__).parent) |
|
|
model_path = os.path.join(path, 'saved_model', 'pfann_fma_m.pt') |
|
|
|
|
|
|
|
|
model_path = os.path.join(path, 'saved_model', 'nnfp_fma.pt') |
|
|
if model_path.endswith('.onnx'): |
|
|
if model_path.endswith('.onnx'): |
|
|
log.warning('Using onnx.') |
|
|
log.warning('Using onnx.') |
|
|
self.model = onnxruntime.InferenceSession( |
|
|
self.model = onnxruntime.InferenceSession( |
|
|