export.py replace check_file -> check_yaml (#8852)
* export.py replace `check_file` -> `check_yaml` * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
1 parent
6884da3a32
commit
628c05ca6f
1 file changed
+2
-2
@@ -67,7 +67,7 @@ if platform.system() != 'Windows':
|
|||||||
from models.experimental import attempt_load
|
from models.experimental import attempt_load
|
||||||
from models.yolo import Detect
|
from models.yolo import Detect
|
||||||
from utils.dataloaders import LoadImages
|
from utils.dataloaders import LoadImages
|
||||||
from utils.general import (LOGGER, check_dataset, check_file, check_img_size, check_requirements, check_version,
|
from utils.general import (LOGGER, check_dataset, check_img_size, check_requirements, check_version, check_yaml,
|
||||||
colorstr, file_size, print_args, url2file)
|
colorstr, file_size, print_args, url2file)
|
||||||
from utils.torch_utils import select_device
|
from utils.torch_utils import select_device
|
||||||
|
|
||||||
@@ -371,7 +371,7 @@ def export_tflite(keras_model, im, file, int8, data, nms, agnostic_nms, prefix=c
|
|||||||
converter.optimizations = [tf.lite.Optimize.DEFAULT]
|
converter.optimizations = [tf.lite.Optimize.DEFAULT]
|
||||||
if int8:
|
if int8:
|
||||||
from models.tf import representative_dataset_gen
|
from models.tf import representative_dataset_gen
|
||||||
dataset = LoadImages(check_dataset(check_file(data))['train'], img_size=imgsz, auto=False)
|
dataset = LoadImages(check_dataset(check_yaml(data))['train'], img_size=imgsz, auto=False)
|
||||||
converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib=100)
|
converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib=100)
|
||||||
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
|
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
|
||||||
converter.target_spec.supported_types = []
|
converter.target_spec.supported_types = []
|
||||||
|
|||||||
Reference in new issue
Block a user