You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

export.py 25KB

Add EdgeTPU support (#3630) * Add models/tf.py for TensorFlow and TFLite export * Set auto=False for int8 calibration * Update requirements.txt for TensorFlow and TFLite export * Read anchors directly from PyTorch weights * Add --tf-nms to append NMS in TensorFlow SavedModel and GraphDef export * Remove check_anchor_order, check_file, set_logging from import * Reformat code and optimize imports * Autodownload model and check cfg * update --source path, img-size to 320, single output * Adjust representative_dataset * Put representative dataset in tfl_int8 block * detect.py TF inference * weights to string * weights to string * cleanup tf.py * Add --dynamic-batch-size * Add xywh normalization to reduce calibration error * Update requirements.txt TensorFlow 2.3.1 -> 2.4.0 to avoid int8 quantization error * Fix imports Move C3 from models.experimental to models.common * Add models/tf.py for TensorFlow and TFLite export * Set auto=False for int8 calibration * Update requirements.txt for TensorFlow and TFLite export * Read anchors directly from PyTorch weights * Add --tf-nms to append NMS in TensorFlow SavedModel and GraphDef export * Remove check_anchor_order, check_file, set_logging from import * Reformat code and optimize imports * Autodownload model and check cfg * update --source path, img-size to 320, single output * Adjust representative_dataset * detect.py TF inference * Put representative dataset in tfl_int8 block * weights to string * weights to string * cleanup tf.py * Add --dynamic-batch-size * Add xywh normalization to reduce calibration error * Update requirements.txt TensorFlow 2.3.1 -> 2.4.0 to avoid int8 quantization error * Fix imports Move C3 from models.experimental to models.common * implement C3() and SiLU() * Add TensorFlow and TFLite Detection * Add --tfl-detect for TFLite Detection * Add int8 quantized TFLite inference in detect.py * Add --edgetpu for Edge TPU detection * Fix --img-size to add rectangle TensorFlow and TFLite input * Add --no-tf-nms to detect objects using models combined with TensorFlow NMS * Fix --img-size list type input * Update README.md * Add Android project for TFLite inference * Upgrade TensorFlow v2.3.1 -> v2.4.0 * Disable normalization of xywh * Rewrite names init in detect.py * Change input resolution 640 -> 320 on Android * Disable NNAPI * Update README.me --img 640 -> 320 * Update README.me for Edge TPU * Update README.md * Fix reshape dim to support dynamic batching * Fix reshape dim to support dynamic batching * Add epsilon argument in tf_BN, which is different between TF and PT * Set stride to None if not using PyTorch, and do not warmup without PyTorch * Add list support in check_img_size() * Add list input support in detect.py * sys.path.append('./') to run from yolov5/ * Add int8 quantization support for TensorFlow 2.5 * Add get_coco128.sh * Remove --no-tfl-detect in models/tf.py (Use tf-android-tfl-detect branch for EdgeTPU) * Update requirements.txt * Replace torch.load() with attempt_load() * Update requirements.txt * Add --tf-raw-resize to set half_pixel_centers=False * Remove android directory * Update README.md * Update README.md * Add multiple OS support for EdgeTPU detection * Fix export and detect * Export 3 YOLO heads with Edge TPU models * Remove xywh denormalization with Edge TPU models in detect.py * Fix saved_model and pb detect error * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix pre-commit.ci failure * Add edgetpu in export.py docstring * Fix Edge TPU model detection exported by TF 2.7 * Add class names for TF/TFLite in DetectMultibackend * Fix assignment with nl in TFLite Detection * Add check when getting Edge TPU compiler version * Add UTF-8 encoding in opening --data file for Windows * Remove redundant TensorFlow import * Add Edge TPU in export.py's docstring * Add the detect layer in Edge TPU model conversion * Default `dnn=False` * Cleanup data.yaml loading * Update detect.py * Update val.py * Comments and generalize data.yaml names Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com> Co-authored-by: unknown <fangjiacong@ut.cn> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2 years ago
Add EdgeTPU support (#3630) * Add models/tf.py for TensorFlow and TFLite export * Set auto=False for int8 calibration * Update requirements.txt for TensorFlow and TFLite export * Read anchors directly from PyTorch weights * Add --tf-nms to append NMS in TensorFlow SavedModel and GraphDef export * Remove check_anchor_order, check_file, set_logging from import * Reformat code and optimize imports * Autodownload model and check cfg * update --source path, img-size to 320, single output * Adjust representative_dataset * Put representative dataset in tfl_int8 block * detect.py TF inference * weights to string * weights to string * cleanup tf.py * Add --dynamic-batch-size * Add xywh normalization to reduce calibration error * Update requirements.txt TensorFlow 2.3.1 -> 2.4.0 to avoid int8 quantization error * Fix imports Move C3 from models.experimental to models.common * Add models/tf.py for TensorFlow and TFLite export * Set auto=False for int8 calibration * Update requirements.txt for TensorFlow and TFLite export * Read anchors directly from PyTorch weights * Add --tf-nms to append NMS in TensorFlow SavedModel and GraphDef export * Remove check_anchor_order, check_file, set_logging from import * Reformat code and optimize imports * Autodownload model and check cfg * update --source path, img-size to 320, single output * Adjust representative_dataset * detect.py TF inference * Put representative dataset in tfl_int8 block * weights to string * weights to string * cleanup tf.py * Add --dynamic-batch-size * Add xywh normalization to reduce calibration error * Update requirements.txt TensorFlow 2.3.1 -> 2.4.0 to avoid int8 quantization error * Fix imports Move C3 from models.experimental to models.common * implement C3() and SiLU() * Add TensorFlow and TFLite Detection * Add --tfl-detect for TFLite Detection * Add int8 quantized TFLite inference in detect.py * Add --edgetpu for Edge TPU detection * Fix --img-size to add rectangle TensorFlow and TFLite input * Add --no-tf-nms to detect objects using models combined with TensorFlow NMS * Fix --img-size list type input * Update README.md * Add Android project for TFLite inference * Upgrade TensorFlow v2.3.1 -> v2.4.0 * Disable normalization of xywh * Rewrite names init in detect.py * Change input resolution 640 -> 320 on Android * Disable NNAPI * Update README.me --img 640 -> 320 * Update README.me for Edge TPU * Update README.md * Fix reshape dim to support dynamic batching * Fix reshape dim to support dynamic batching * Add epsilon argument in tf_BN, which is different between TF and PT * Set stride to None if not using PyTorch, and do not warmup without PyTorch * Add list support in check_img_size() * Add list input support in detect.py * sys.path.append('./') to run from yolov5/ * Add int8 quantization support for TensorFlow 2.5 * Add get_coco128.sh * Remove --no-tfl-detect in models/tf.py (Use tf-android-tfl-detect branch for EdgeTPU) * Update requirements.txt * Replace torch.load() with attempt_load() * Update requirements.txt * Add --tf-raw-resize to set half_pixel_centers=False * Remove android directory * Update README.md * Update README.md * Add multiple OS support for EdgeTPU detection * Fix export and detect * Export 3 YOLO heads with Edge TPU models * Remove xywh denormalization with Edge TPU models in detect.py * Fix saved_model and pb detect error * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix pre-commit.ci failure * Add edgetpu in export.py docstring * Fix Edge TPU model detection exported by TF 2.7 * Add class names for TF/TFLite in DetectMultibackend * Fix assignment with nl in TFLite Detection * Add check when getting Edge TPU compiler version * Add UTF-8 encoding in opening --data file for Windows * Remove redundant TensorFlow import * Add Edge TPU in export.py's docstring * Add the detect layer in Edge TPU model conversion * Default `dnn=False` * Cleanup data.yaml loading * Update detect.py * Update val.py * Comments and generalize data.yaml names Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com> Co-authored-by: unknown <fangjiacong@ut.cn> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2 years ago
Add EdgeTPU support (#3630) * Add models/tf.py for TensorFlow and TFLite export * Set auto=False for int8 calibration * Update requirements.txt for TensorFlow and TFLite export * Read anchors directly from PyTorch weights * Add --tf-nms to append NMS in TensorFlow SavedModel and GraphDef export * Remove check_anchor_order, check_file, set_logging from import * Reformat code and optimize imports * Autodownload model and check cfg * update --source path, img-size to 320, single output * Adjust representative_dataset * Put representative dataset in tfl_int8 block * detect.py TF inference * weights to string * weights to string * cleanup tf.py * Add --dynamic-batch-size * Add xywh normalization to reduce calibration error * Update requirements.txt TensorFlow 2.3.1 -> 2.4.0 to avoid int8 quantization error * Fix imports Move C3 from models.experimental to models.common * Add models/tf.py for TensorFlow and TFLite export * Set auto=False for int8 calibration * Update requirements.txt for TensorFlow and TFLite export * Read anchors directly from PyTorch weights * Add --tf-nms to append NMS in TensorFlow SavedModel and GraphDef export * Remove check_anchor_order, check_file, set_logging from import * Reformat code and optimize imports * Autodownload model and check cfg * update --source path, img-size to 320, single output * Adjust representative_dataset * detect.py TF inference * Put representative dataset in tfl_int8 block * weights to string * weights to string * cleanup tf.py * Add --dynamic-batch-size * Add xywh normalization to reduce calibration error * Update requirements.txt TensorFlow 2.3.1 -> 2.4.0 to avoid int8 quantization error * Fix imports Move C3 from models.experimental to models.common * implement C3() and SiLU() * Add TensorFlow and TFLite Detection * Add --tfl-detect for TFLite Detection * Add int8 quantized TFLite inference in detect.py * Add --edgetpu for Edge TPU detection * Fix --img-size to add rectangle TensorFlow and TFLite input * Add --no-tf-nms to detect objects using models combined with TensorFlow NMS * Fix --img-size list type input * Update README.md * Add Android project for TFLite inference * Upgrade TensorFlow v2.3.1 -> v2.4.0 * Disable normalization of xywh * Rewrite names init in detect.py * Change input resolution 640 -> 320 on Android * Disable NNAPI * Update README.me --img 640 -> 320 * Update README.me for Edge TPU * Update README.md * Fix reshape dim to support dynamic batching * Fix reshape dim to support dynamic batching * Add epsilon argument in tf_BN, which is different between TF and PT * Set stride to None if not using PyTorch, and do not warmup without PyTorch * Add list support in check_img_size() * Add list input support in detect.py * sys.path.append('./') to run from yolov5/ * Add int8 quantization support for TensorFlow 2.5 * Add get_coco128.sh * Remove --no-tfl-detect in models/tf.py (Use tf-android-tfl-detect branch for EdgeTPU) * Update requirements.txt * Replace torch.load() with attempt_load() * Update requirements.txt * Add --tf-raw-resize to set half_pixel_centers=False * Remove android directory * Update README.md * Update README.md * Add multiple OS support for EdgeTPU detection * Fix export and detect * Export 3 YOLO heads with Edge TPU models * Remove xywh denormalization with Edge TPU models in detect.py * Fix saved_model and pb detect error * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix pre-commit.ci failure * Add edgetpu in export.py docstring * Fix Edge TPU model detection exported by TF 2.7 * Add class names for TF/TFLite in DetectMultibackend * Fix assignment with nl in TFLite Detection * Add check when getting Edge TPU compiler version * Add UTF-8 encoding in opening --data file for Windows * Remove redundant TensorFlow import * Add Edge TPU in export.py's docstring * Add the detect layer in Edge TPU model conversion * Default `dnn=False` * Cleanup data.yaml loading * Update detect.py * Update val.py * Comments and generalize data.yaml names Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com> Co-authored-by: unknown <fangjiacong@ut.cn> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2 years ago
Add EdgeTPU support (#3630) * Add models/tf.py for TensorFlow and TFLite export * Set auto=False for int8 calibration * Update requirements.txt for TensorFlow and TFLite export * Read anchors directly from PyTorch weights * Add --tf-nms to append NMS in TensorFlow SavedModel and GraphDef export * Remove check_anchor_order, check_file, set_logging from import * Reformat code and optimize imports * Autodownload model and check cfg * update --source path, img-size to 320, single output * Adjust representative_dataset * Put representative dataset in tfl_int8 block * detect.py TF inference * weights to string * weights to string * cleanup tf.py * Add --dynamic-batch-size * Add xywh normalization to reduce calibration error * Update requirements.txt TensorFlow 2.3.1 -> 2.4.0 to avoid int8 quantization error * Fix imports Move C3 from models.experimental to models.common * Add models/tf.py for TensorFlow and TFLite export * Set auto=False for int8 calibration * Update requirements.txt for TensorFlow and TFLite export * Read anchors directly from PyTorch weights * Add --tf-nms to append NMS in TensorFlow SavedModel and GraphDef export * Remove check_anchor_order, check_file, set_logging from import * Reformat code and optimize imports * Autodownload model and check cfg * update --source path, img-size to 320, single output * Adjust representative_dataset * detect.py TF inference * Put representative dataset in tfl_int8 block * weights to string * weights to string * cleanup tf.py * Add --dynamic-batch-size * Add xywh normalization to reduce calibration error * Update requirements.txt TensorFlow 2.3.1 -> 2.4.0 to avoid int8 quantization error * Fix imports Move C3 from models.experimental to models.common * implement C3() and SiLU() * Add TensorFlow and TFLite Detection * Add --tfl-detect for TFLite Detection * Add int8 quantized TFLite inference in detect.py * Add --edgetpu for Edge TPU detection * Fix --img-size to add rectangle TensorFlow and TFLite input * Add --no-tf-nms to detect objects using models combined with TensorFlow NMS * Fix --img-size list type input * Update README.md * Add Android project for TFLite inference * Upgrade TensorFlow v2.3.1 -> v2.4.0 * Disable normalization of xywh * Rewrite names init in detect.py * Change input resolution 640 -> 320 on Android * Disable NNAPI * Update README.me --img 640 -> 320 * Update README.me for Edge TPU * Update README.md * Fix reshape dim to support dynamic batching * Fix reshape dim to support dynamic batching * Add epsilon argument in tf_BN, which is different between TF and PT * Set stride to None if not using PyTorch, and do not warmup without PyTorch * Add list support in check_img_size() * Add list input support in detect.py * sys.path.append('./') to run from yolov5/ * Add int8 quantization support for TensorFlow 2.5 * Add get_coco128.sh * Remove --no-tfl-detect in models/tf.py (Use tf-android-tfl-detect branch for EdgeTPU) * Update requirements.txt * Replace torch.load() with attempt_load() * Update requirements.txt * Add --tf-raw-resize to set half_pixel_centers=False * Remove android directory * Update README.md * Update README.md * Add multiple OS support for EdgeTPU detection * Fix export and detect * Export 3 YOLO heads with Edge TPU models * Remove xywh denormalization with Edge TPU models in detect.py * Fix saved_model and pb detect error * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix pre-commit.ci failure * Add edgetpu in export.py docstring * Fix Edge TPU model detection exported by TF 2.7 * Add class names for TF/TFLite in DetectMultibackend * Fix assignment with nl in TFLite Detection * Add check when getting Edge TPU compiler version * Add UTF-8 encoding in opening --data file for Windows * Remove redundant TensorFlow import * Add Edge TPU in export.py's docstring * Add the detect layer in Edge TPU model conversion * Default `dnn=False` * Cleanup data.yaml loading * Update detect.py * Update val.py * Comments and generalize data.yaml names Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com> Co-authored-by: unknown <fangjiacong@ut.cn> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2 years ago
Add EdgeTPU support (#3630) * Add models/tf.py for TensorFlow and TFLite export * Set auto=False for int8 calibration * Update requirements.txt for TensorFlow and TFLite export * Read anchors directly from PyTorch weights * Add --tf-nms to append NMS in TensorFlow SavedModel and GraphDef export * Remove check_anchor_order, check_file, set_logging from import * Reformat code and optimize imports * Autodownload model and check cfg * update --source path, img-size to 320, single output * Adjust representative_dataset * Put representative dataset in tfl_int8 block * detect.py TF inference * weights to string * weights to string * cleanup tf.py * Add --dynamic-batch-size * Add xywh normalization to reduce calibration error * Update requirements.txt TensorFlow 2.3.1 -> 2.4.0 to avoid int8 quantization error * Fix imports Move C3 from models.experimental to models.common * Add models/tf.py for TensorFlow and TFLite export * Set auto=False for int8 calibration * Update requirements.txt for TensorFlow and TFLite export * Read anchors directly from PyTorch weights * Add --tf-nms to append NMS in TensorFlow SavedModel and GraphDef export * Remove check_anchor_order, check_file, set_logging from import * Reformat code and optimize imports * Autodownload model and check cfg * update --source path, img-size to 320, single output * Adjust representative_dataset * detect.py TF inference * Put representative dataset in tfl_int8 block * weights to string * weights to string * cleanup tf.py * Add --dynamic-batch-size * Add xywh normalization to reduce calibration error * Update requirements.txt TensorFlow 2.3.1 -> 2.4.0 to avoid int8 quantization error * Fix imports Move C3 from models.experimental to models.common * implement C3() and SiLU() * Add TensorFlow and TFLite Detection * Add --tfl-detect for TFLite Detection * Add int8 quantized TFLite inference in detect.py * Add --edgetpu for Edge TPU detection * Fix --img-size to add rectangle TensorFlow and TFLite input * Add --no-tf-nms to detect objects using models combined with TensorFlow NMS * Fix --img-size list type input * Update README.md * Add Android project for TFLite inference * Upgrade TensorFlow v2.3.1 -> v2.4.0 * Disable normalization of xywh * Rewrite names init in detect.py * Change input resolution 640 -> 320 on Android * Disable NNAPI * Update README.me --img 640 -> 320 * Update README.me for Edge TPU * Update README.md * Fix reshape dim to support dynamic batching * Fix reshape dim to support dynamic batching * Add epsilon argument in tf_BN, which is different between TF and PT * Set stride to None if not using PyTorch, and do not warmup without PyTorch * Add list support in check_img_size() * Add list input support in detect.py * sys.path.append('./') to run from yolov5/ * Add int8 quantization support for TensorFlow 2.5 * Add get_coco128.sh * Remove --no-tfl-detect in models/tf.py (Use tf-android-tfl-detect branch for EdgeTPU) * Update requirements.txt * Replace torch.load() with attempt_load() * Update requirements.txt * Add --tf-raw-resize to set half_pixel_centers=False * Remove android directory * Update README.md * Update README.md * Add multiple OS support for EdgeTPU detection * Fix export and detect * Export 3 YOLO heads with Edge TPU models * Remove xywh denormalization with Edge TPU models in detect.py * Fix saved_model and pb detect error * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix pre-commit.ci failure * Add edgetpu in export.py docstring * Fix Edge TPU model detection exported by TF 2.7 * Add class names for TF/TFLite in DetectMultibackend * Fix assignment with nl in TFLite Detection * Add check when getting Edge TPU compiler version * Add UTF-8 encoding in opening --data file for Windows * Remove redundant TensorFlow import * Add Edge TPU in export.py's docstring * Add the detect layer in Edge TPU model conversion * Default `dnn=False` * Cleanup data.yaml loading * Update detect.py * Update val.py * Comments and generalize data.yaml names Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com> Co-authored-by: unknown <fangjiacong@ut.cn> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2 years ago
Merge `develop` branch into `master` (#3518) * update ci-testing.yml (#3322) * update ci-testing.yml * update greetings.yml * bring back os matrix * update ci-testing.yml (#3322) * update ci-testing.yml * update greetings.yml * bring back os matrix * Enable direct `--weights URL` definition (#3373) * Enable direct `--weights URL` definition @KalenMike this PR will enable direct --weights URL definition. Example use case: ``` python train.py --weights https://storage.googleapis.com/bucket/dir/model.pt ``` * cleanup * bug fixes * weights = attempt_download(weights) * Update experimental.py * Update hubconf.py * return bug fix * comment mirror * min_bytes * Update tutorial.ipynb (#3368) add Open in Kaggle badge * `cv2.imread(img, -1)` for IMREAD_UNCHANGED (#3379) * Update datasets.py * comment Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com> * COCO evolution fix (#3388) * COCO evolution fix * cleanup * update print * print fix * Create `is_pip()` function (#3391) Returns `True` if file is part of pip package. Useful for contextual behavior modification. ```python def is_pip(): # Is file in a pip package? return 'site-packages' in Path(__file__).absolute().parts ``` * Revert "`cv2.imread(img, -1)` for IMREAD_UNCHANGED (#3379)" (#3395) This reverts commit 21a9607e00f1365b21d8c4bd81bdbf5fc0efea24. * Update FLOPs description (#3422) * Update README.md * Changing FLOPS to FLOPs. Co-authored-by: BuildTools <unconfigured@null.spigotmc.org> * Parse URL authentication (#3424) * Parse URL authentication * urllib.parse.unquote() * improved error handling * improved error handling * remove %3F * update check_file() * Add FLOPs title to table (#3453) * Suppress jit trace warning + graph once (#3454) * Suppress jit trace warning + graph once Suppress harmless jit trace warning on TensorBoard add_graph call. Also fix multiple add_graph() calls bug, now only on batch 0. * Update train.py * Update MixUp augmentation `alpha=beta=32.0` (#3455) Per VOC empirical results https://github.com/ultralytics/yolov5/issues/3380#issuecomment-853001307 by @developer0hye * Add `timeout()` class (#3460) * Add `timeout()` class * rearrange order * Faster HSV augmentation (#3462) remove datatype conversion process that can be skipped * Add `check_git_status()` 5 second timeout (#3464) * Add check_git_status() 5 second timeout This should prevent the SSH Git bug that we were discussing @KalenMike * cleanup * replace timeout with check_output built-in timeout * Improved `check_requirements()` offline-handling (#3466) Improve robustness of `check_requirements()` function to offline environments (do not attempt pip installs when offline). * Add `output_names` argument for ONNX export with dynamic axes (#3456) * Add output names & dynamic axes for onnx export Add output_names and dynamic_axes names for all outputs in torch.onnx.export. The first four outputs of the model will have names output0, output1, output2, output3 * use first output only + cleanup Co-authored-by: Samridha Shrestha <samridha.shrestha@g42.ai> Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com> * Revert FP16 `test.py` and `detect.py` inference to FP32 default (#3423) * fixed inference bug ,while use half precision * replace --use-half with --half * replace space and PEP8 in detect.py * PEP8 detect.py * update --half help comment * Update test.py * revert space Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com> * Add additional links/resources to stale.yml message (#3467) * Update stale.yml * cleanup * Update stale.yml * reformat * Update stale.yml HUB URL (#3468) * Stale `github.actor` bug fix (#3483) * Explicit `model.eval()` call `if opt.train=False` (#3475) * call model.eval() when opt.train is False call model.eval() when opt.train is False * single-line if statement * cleanup Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com> * check_requirements() exclude `opencv-python` (#3495) Fix for 3rd party or contrib versions of installed OpenCV as in https://github.com/ultralytics/yolov5/issues/3494. * Earlier `assert` for cpu and half option (#3508) * early assert for cpu and half option early assert for cpu and half option * Modified comment Modified comment * Update tutorial.ipynb (#3510) * Reduce test.py results spacing (#3511) * Update README.md (#3512) * Update README.md Minor modifications * 850 width * Update greetings.yml revert greeting change as PRs will now merge to master. Co-authored-by: Piotr Skalski <SkalskiP@users.noreply.github.com> Co-authored-by: SkalskiP <piotr.skalski92@gmail.com> Co-authored-by: Peretz Cohen <pizzaz93@users.noreply.github.com> Co-authored-by: tudoulei <34886368+tudoulei@users.noreply.github.com> Co-authored-by: chocosaj <chocosaj@users.noreply.github.com> Co-authored-by: BuildTools <unconfigured@null.spigotmc.org> Co-authored-by: Yonghye Kwon <developer.0hye@gmail.com> Co-authored-by: Sam_S <SamSamhuns@users.noreply.github.com> Co-authored-by: Samridha Shrestha <samridha.shrestha@g42.ai> Co-authored-by: edificewang <609552430@qq.com>
3 years ago
Add EdgeTPU support (#3630) * Add models/tf.py for TensorFlow and TFLite export * Set auto=False for int8 calibration * Update requirements.txt for TensorFlow and TFLite export * Read anchors directly from PyTorch weights * Add --tf-nms to append NMS in TensorFlow SavedModel and GraphDef export * Remove check_anchor_order, check_file, set_logging from import * Reformat code and optimize imports * Autodownload model and check cfg * update --source path, img-size to 320, single output * Adjust representative_dataset * Put representative dataset in tfl_int8 block * detect.py TF inference * weights to string * weights to string * cleanup tf.py * Add --dynamic-batch-size * Add xywh normalization to reduce calibration error * Update requirements.txt TensorFlow 2.3.1 -> 2.4.0 to avoid int8 quantization error * Fix imports Move C3 from models.experimental to models.common * Add models/tf.py for TensorFlow and TFLite export * Set auto=False for int8 calibration * Update requirements.txt for TensorFlow and TFLite export * Read anchors directly from PyTorch weights * Add --tf-nms to append NMS in TensorFlow SavedModel and GraphDef export * Remove check_anchor_order, check_file, set_logging from import * Reformat code and optimize imports * Autodownload model and check cfg * update --source path, img-size to 320, single output * Adjust representative_dataset * detect.py TF inference * Put representative dataset in tfl_int8 block * weights to string * weights to string * cleanup tf.py * Add --dynamic-batch-size * Add xywh normalization to reduce calibration error * Update requirements.txt TensorFlow 2.3.1 -> 2.4.0 to avoid int8 quantization error * Fix imports Move C3 from models.experimental to models.common * implement C3() and SiLU() * Add TensorFlow and TFLite Detection * Add --tfl-detect for TFLite Detection * Add int8 quantized TFLite inference in detect.py * Add --edgetpu for Edge TPU detection * Fix --img-size to add rectangle TensorFlow and TFLite input * Add --no-tf-nms to detect objects using models combined with TensorFlow NMS * Fix --img-size list type input * Update README.md * Add Android project for TFLite inference * Upgrade TensorFlow v2.3.1 -> v2.4.0 * Disable normalization of xywh * Rewrite names init in detect.py * Change input resolution 640 -> 320 on Android * Disable NNAPI * Update README.me --img 640 -> 320 * Update README.me for Edge TPU * Update README.md * Fix reshape dim to support dynamic batching * Fix reshape dim to support dynamic batching * Add epsilon argument in tf_BN, which is different between TF and PT * Set stride to None if not using PyTorch, and do not warmup without PyTorch * Add list support in check_img_size() * Add list input support in detect.py * sys.path.append('./') to run from yolov5/ * Add int8 quantization support for TensorFlow 2.5 * Add get_coco128.sh * Remove --no-tfl-detect in models/tf.py (Use tf-android-tfl-detect branch for EdgeTPU) * Update requirements.txt * Replace torch.load() with attempt_load() * Update requirements.txt * Add --tf-raw-resize to set half_pixel_centers=False * Remove android directory * Update README.md * Update README.md * Add multiple OS support for EdgeTPU detection * Fix export and detect * Export 3 YOLO heads with Edge TPU models * Remove xywh denormalization with Edge TPU models in detect.py * Fix saved_model and pb detect error * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix pre-commit.ci failure * Add edgetpu in export.py docstring * Fix Edge TPU model detection exported by TF 2.7 * Add class names for TF/TFLite in DetectMultibackend * Fix assignment with nl in TFLite Detection * Add check when getting Edge TPU compiler version * Add UTF-8 encoding in opening --data file for Windows * Remove redundant TensorFlow import * Add Edge TPU in export.py's docstring * Add the detect layer in Edge TPU model conversion * Default `dnn=False` * Cleanup data.yaml loading * Update detect.py * Update val.py * Comments and generalize data.yaml names Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com> Co-authored-by: unknown <fangjiacong@ut.cn> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2 years ago
Add EdgeTPU support (#3630) * Add models/tf.py for TensorFlow and TFLite export * Set auto=False for int8 calibration * Update requirements.txt for TensorFlow and TFLite export * Read anchors directly from PyTorch weights * Add --tf-nms to append NMS in TensorFlow SavedModel and GraphDef export * Remove check_anchor_order, check_file, set_logging from import * Reformat code and optimize imports * Autodownload model and check cfg * update --source path, img-size to 320, single output * Adjust representative_dataset * Put representative dataset in tfl_int8 block * detect.py TF inference * weights to string * weights to string * cleanup tf.py * Add --dynamic-batch-size * Add xywh normalization to reduce calibration error * Update requirements.txt TensorFlow 2.3.1 -> 2.4.0 to avoid int8 quantization error * Fix imports Move C3 from models.experimental to models.common * Add models/tf.py for TensorFlow and TFLite export * Set auto=False for int8 calibration * Update requirements.txt for TensorFlow and TFLite export * Read anchors directly from PyTorch weights * Add --tf-nms to append NMS in TensorFlow SavedModel and GraphDef export * Remove check_anchor_order, check_file, set_logging from import * Reformat code and optimize imports * Autodownload model and check cfg * update --source path, img-size to 320, single output * Adjust representative_dataset * detect.py TF inference * Put representative dataset in tfl_int8 block * weights to string * weights to string * cleanup tf.py * Add --dynamic-batch-size * Add xywh normalization to reduce calibration error * Update requirements.txt TensorFlow 2.3.1 -> 2.4.0 to avoid int8 quantization error * Fix imports Move C3 from models.experimental to models.common * implement C3() and SiLU() * Add TensorFlow and TFLite Detection * Add --tfl-detect for TFLite Detection * Add int8 quantized TFLite inference in detect.py * Add --edgetpu for Edge TPU detection * Fix --img-size to add rectangle TensorFlow and TFLite input * Add --no-tf-nms to detect objects using models combined with TensorFlow NMS * Fix --img-size list type input * Update README.md * Add Android project for TFLite inference * Upgrade TensorFlow v2.3.1 -> v2.4.0 * Disable normalization of xywh * Rewrite names init in detect.py * Change input resolution 640 -> 320 on Android * Disable NNAPI * Update README.me --img 640 -> 320 * Update README.me for Edge TPU * Update README.md * Fix reshape dim to support dynamic batching * Fix reshape dim to support dynamic batching * Add epsilon argument in tf_BN, which is different between TF and PT * Set stride to None if not using PyTorch, and do not warmup without PyTorch * Add list support in check_img_size() * Add list input support in detect.py * sys.path.append('./') to run from yolov5/ * Add int8 quantization support for TensorFlow 2.5 * Add get_coco128.sh * Remove --no-tfl-detect in models/tf.py (Use tf-android-tfl-detect branch for EdgeTPU) * Update requirements.txt * Replace torch.load() with attempt_load() * Update requirements.txt * Add --tf-raw-resize to set half_pixel_centers=False * Remove android directory * Update README.md * Update README.md * Add multiple OS support for EdgeTPU detection * Fix export and detect * Export 3 YOLO heads with Edge TPU models * Remove xywh denormalization with Edge TPU models in detect.py * Fix saved_model and pb detect error * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix pre-commit.ci failure * Add edgetpu in export.py docstring * Fix Edge TPU model detection exported by TF 2.7 * Add class names for TF/TFLite in DetectMultibackend * Fix assignment with nl in TFLite Detection * Add check when getting Edge TPU compiler version * Add UTF-8 encoding in opening --data file for Windows * Remove redundant TensorFlow import * Add Edge TPU in export.py's docstring * Add the detect layer in Edge TPU model conversion * Default `dnn=False` * Cleanup data.yaml loading * Update detect.py * Update val.py * Comments and generalize data.yaml names Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com> Co-authored-by: unknown <fangjiacong@ut.cn> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2 years ago
Add EdgeTPU support (#3630) * Add models/tf.py for TensorFlow and TFLite export * Set auto=False for int8 calibration * Update requirements.txt for TensorFlow and TFLite export * Read anchors directly from PyTorch weights * Add --tf-nms to append NMS in TensorFlow SavedModel and GraphDef export * Remove check_anchor_order, check_file, set_logging from import * Reformat code and optimize imports * Autodownload model and check cfg * update --source path, img-size to 320, single output * Adjust representative_dataset * Put representative dataset in tfl_int8 block * detect.py TF inference * weights to string * weights to string * cleanup tf.py * Add --dynamic-batch-size * Add xywh normalization to reduce calibration error * Update requirements.txt TensorFlow 2.3.1 -> 2.4.0 to avoid int8 quantization error * Fix imports Move C3 from models.experimental to models.common * Add models/tf.py for TensorFlow and TFLite export * Set auto=False for int8 calibration * Update requirements.txt for TensorFlow and TFLite export * Read anchors directly from PyTorch weights * Add --tf-nms to append NMS in TensorFlow SavedModel and GraphDef export * Remove check_anchor_order, check_file, set_logging from import * Reformat code and optimize imports * Autodownload model and check cfg * update --source path, img-size to 320, single output * Adjust representative_dataset * detect.py TF inference * Put representative dataset in tfl_int8 block * weights to string * weights to string * cleanup tf.py * Add --dynamic-batch-size * Add xywh normalization to reduce calibration error * Update requirements.txt TensorFlow 2.3.1 -> 2.4.0 to avoid int8 quantization error * Fix imports Move C3 from models.experimental to models.common * implement C3() and SiLU() * Add TensorFlow and TFLite Detection * Add --tfl-detect for TFLite Detection * Add int8 quantized TFLite inference in detect.py * Add --edgetpu for Edge TPU detection * Fix --img-size to add rectangle TensorFlow and TFLite input * Add --no-tf-nms to detect objects using models combined with TensorFlow NMS * Fix --img-size list type input * Update README.md * Add Android project for TFLite inference * Upgrade TensorFlow v2.3.1 -> v2.4.0 * Disable normalization of xywh * Rewrite names init in detect.py * Change input resolution 640 -> 320 on Android * Disable NNAPI * Update README.me --img 640 -> 320 * Update README.me for Edge TPU * Update README.md * Fix reshape dim to support dynamic batching * Fix reshape dim to support dynamic batching * Add epsilon argument in tf_BN, which is different between TF and PT * Set stride to None if not using PyTorch, and do not warmup without PyTorch * Add list support in check_img_size() * Add list input support in detect.py * sys.path.append('./') to run from yolov5/ * Add int8 quantization support for TensorFlow 2.5 * Add get_coco128.sh * Remove --no-tfl-detect in models/tf.py (Use tf-android-tfl-detect branch for EdgeTPU) * Update requirements.txt * Replace torch.load() with attempt_load() * Update requirements.txt * Add --tf-raw-resize to set half_pixel_centers=False * Remove android directory * Update README.md * Update README.md * Add multiple OS support for EdgeTPU detection * Fix export and detect * Export 3 YOLO heads with Edge TPU models * Remove xywh denormalization with Edge TPU models in detect.py * Fix saved_model and pb detect error * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix pre-commit.ci failure * Add edgetpu in export.py docstring * Fix Edge TPU model detection exported by TF 2.7 * Add class names for TF/TFLite in DetectMultibackend * Fix assignment with nl in TFLite Detection * Add check when getting Edge TPU compiler version * Add UTF-8 encoding in opening --data file for Windows * Remove redundant TensorFlow import * Add Edge TPU in export.py's docstring * Add the detect layer in Edge TPU model conversion * Default `dnn=False` * Cleanup data.yaml loading * Update detect.py * Update val.py * Comments and generalize data.yaml names Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com> Co-authored-by: unknown <fangjiacong@ut.cn> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2 years ago
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523
  1. # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
  2. """
  3. Export a YOLOv5 PyTorch model to other formats. TensorFlow exports authored by https://github.com/zldrobit
  4. Format | `export.py --include` | Model
  5. --- | --- | ---
  6. PyTorch | - | yolov5s.pt
  7. TorchScript | `torchscript` | yolov5s.torchscript
  8. ONNX | `onnx` | yolov5s.onnx
  9. OpenVINO | `openvino` | yolov5s_openvino_model/
  10. TensorRT | `engine` | yolov5s.engine
  11. CoreML | `coreml` | yolov5s.mlmodel
  12. TensorFlow SavedModel | `saved_model` | yolov5s_saved_model/
  13. TensorFlow GraphDef | `pb` | yolov5s.pb
  14. TensorFlow Lite | `tflite` | yolov5s.tflite
  15. TensorFlow Edge TPU | `edgetpu` | yolov5s_edgetpu.tflite
  16. TensorFlow.js | `tfjs` | yolov5s_web_model/
  17. Requirements:
  18. $ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime openvino-dev tensorflow-cpu # CPU
  19. $ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime-gpu openvino-dev tensorflow # GPU
  20. Usage:
  21. $ python path/to/export.py --weights yolov5s.pt --include torchscript onnx openvino engine coreml tflite ...
  22. Inference:
  23. $ python path/to/detect.py --weights yolov5s.pt # PyTorch
  24. yolov5s.torchscript # TorchScript
  25. yolov5s.onnx # ONNX Runtime or OpenCV DNN with --dnn
  26. yolov5s.xml # OpenVINO
  27. yolov5s.engine # TensorRT
  28. yolov5s.mlmodel # CoreML (MacOS-only)
  29. yolov5s_saved_model # TensorFlow SavedModel
  30. yolov5s.pb # TensorFlow GraphDef
  31. yolov5s.tflite # TensorFlow Lite
  32. yolov5s_edgetpu.tflite # TensorFlow Edge TPU
  33. TensorFlow.js:
  34. $ cd .. && git clone https://github.com/zldrobit/tfjs-yolov5-example.git && cd tfjs-yolov5-example
  35. $ npm install
  36. $ ln -s ../../yolov5/yolov5s_web_model public/yolov5s_web_model
  37. $ npm start
  38. """
  39. import argparse
  40. import json
  41. import os
  42. import platform
  43. import subprocess
  44. import sys
  45. import time
  46. import warnings
  47. from pathlib import Path
  48. import torch
  49. import torch.nn as nn
  50. from torch.utils.mobile_optimizer import optimize_for_mobile
  51. FILE = Path(__file__).resolve()
  52. ROOT = FILE.parents[0] # YOLOv5 root directory
  53. if str(ROOT) not in sys.path:
  54. sys.path.append(str(ROOT)) # add ROOT to PATH
  55. ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
  56. from models.common import Conv
  57. from models.experimental import attempt_load
  58. from models.yolo import Detect
  59. from utils.activations import SiLU
  60. from utils.datasets import LoadImages
  61. from utils.general import (LOGGER, check_dataset, check_img_size, check_requirements, check_version, colorstr,
  62. file_size, print_args, url2file)
  63. from utils.torch_utils import select_device
  64. def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:')):
  65. # YOLOv5 TorchScript model export
  66. try:
  67. LOGGER.info(f'\n{prefix} starting export with torch {torch.__version__}...')
  68. f = file.with_suffix('.torchscript')
  69. ts = torch.jit.trace(model, im, strict=False)
  70. d = {"shape": im.shape, "stride": int(max(model.stride)), "names": model.names}
  71. extra_files = {'config.txt': json.dumps(d)} # torch._C.ExtraFilesMap()
  72. if optimize: # https://pytorch.org/tutorials/recipes/mobile_interpreter.html
  73. optimize_for_mobile(ts)._save_for_lite_interpreter(str(f), _extra_files=extra_files)
  74. else:
  75. ts.save(str(f), _extra_files=extra_files)
  76. LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  77. return f
  78. except Exception as e:
  79. LOGGER.info(f'{prefix} export failure: {e}')
  80. def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorstr('ONNX:')):
  81. # YOLOv5 ONNX export
  82. try:
  83. check_requirements(('onnx',))
  84. import onnx
  85. LOGGER.info(f'\n{prefix} starting export with onnx {onnx.__version__}...')
  86. f = file.with_suffix('.onnx')
  87. torch.onnx.export(model, im, f, verbose=False, opset_version=opset,
  88. training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,
  89. do_constant_folding=not train,
  90. input_names=['images'],
  91. output_names=['output'],
  92. dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}, # shape(1,3,640,640)
  93. 'output': {0: 'batch', 1: 'anchors'} # shape(1,25200,85)
  94. } if dynamic else None)
  95. # Checks
  96. model_onnx = onnx.load(f) # load onnx model
  97. onnx.checker.check_model(model_onnx) # check onnx model
  98. # LOGGER.info(onnx.helper.printable_graph(model_onnx.graph)) # print
  99. # Simplify
  100. if simplify:
  101. try:
  102. check_requirements(('onnx-simplifier',))
  103. import onnxsim
  104. LOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
  105. model_onnx, check = onnxsim.simplify(
  106. model_onnx,
  107. dynamic_input_shape=dynamic,
  108. input_shapes={'images': list(im.shape)} if dynamic else None)
  109. assert check, 'assert check failed'
  110. onnx.save(model_onnx, f)
  111. except Exception as e:
  112. LOGGER.info(f'{prefix} simplifier failure: {e}')
  113. LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  114. return f
  115. except Exception as e:
  116. LOGGER.info(f'{prefix} export failure: {e}')
  117. def export_openvino(model, im, file, prefix=colorstr('OpenVINO:')):
  118. # YOLOv5 OpenVINO export
  119. try:
  120. check_requirements(('openvino-dev',)) # requires openvino-dev: https://pypi.org/project/openvino-dev/
  121. import openvino.inference_engine as ie
  122. LOGGER.info(f'\n{prefix} starting export with openvino {ie.__version__}...')
  123. f = str(file).replace('.pt', '_openvino_model' + os.sep)
  124. cmd = f"mo --input_model {file.with_suffix('.onnx')} --output_dir {f}"
  125. subprocess.check_output(cmd, shell=True)
  126. LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  127. return f
  128. except Exception as e:
  129. LOGGER.info(f'\n{prefix} export failure: {e}')
  130. def export_coreml(model, im, file, prefix=colorstr('CoreML:')):
  131. # YOLOv5 CoreML export
  132. try:
  133. check_requirements(('coremltools',))
  134. import coremltools as ct
  135. LOGGER.info(f'\n{prefix} starting export with coremltools {ct.__version__}...')
  136. f = file.with_suffix('.mlmodel')
  137. ts = torch.jit.trace(model, im, strict=False) # TorchScript model
  138. ct_model = ct.convert(ts, inputs=[ct.ImageType('image', shape=im.shape, scale=1 / 255, bias=[0, 0, 0])])
  139. ct_model.save(f)
  140. LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  141. return ct_model, f
  142. except Exception as e:
  143. LOGGER.info(f'\n{prefix} export failure: {e}')
  144. return None, None
  145. def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=False, prefix=colorstr('TensorRT:')):
  146. # YOLOv5 TensorRT export https://developer.nvidia.com/tensorrt
  147. try:
  148. check_requirements(('tensorrt',))
  149. import tensorrt as trt
  150. if trt.__version__[0] == '7': # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012
  151. grid = model.model[-1].anchor_grid
  152. model.model[-1].anchor_grid = [a[..., :1, :1, :] for a in grid]
  153. export_onnx(model, im, file, 12, train, False, simplify) # opset 12
  154. model.model[-1].anchor_grid = grid
  155. else: # TensorRT >= 8
  156. check_version(trt.__version__, '8.0.0', hard=True) # require tensorrt>=8.0.0
  157. export_onnx(model, im, file, 13, train, False, simplify) # opset 13
  158. onnx = file.with_suffix('.onnx')
  159. LOGGER.info(f'\n{prefix} starting export with TensorRT {trt.__version__}...')
  160. assert im.device.type != 'cpu', 'export running on CPU but must be on GPU, i.e. `python export.py --device 0`'
  161. assert onnx.exists(), f'failed to export ONNX file: {onnx}'
  162. f = file.with_suffix('.engine') # TensorRT engine file
  163. logger = trt.Logger(trt.Logger.INFO)
  164. if verbose:
  165. logger.min_severity = trt.Logger.Severity.VERBOSE
  166. builder = trt.Builder(logger)
  167. config = builder.create_builder_config()
  168. config.max_workspace_size = workspace * 1 << 30
  169. flag = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
  170. network = builder.create_network(flag)
  171. parser = trt.OnnxParser(network, logger)
  172. if not parser.parse_from_file(str(onnx)):
  173. raise RuntimeError(f'failed to load ONNX file: {onnx}')
  174. inputs = [network.get_input(i) for i in range(network.num_inputs)]
  175. outputs = [network.get_output(i) for i in range(network.num_outputs)]
  176. LOGGER.info(f'{prefix} Network Description:')
  177. for inp in inputs:
  178. LOGGER.info(f'{prefix}\tinput "{inp.name}" with shape {inp.shape} and dtype {inp.dtype}')
  179. for out in outputs:
  180. LOGGER.info(f'{prefix}\toutput "{out.name}" with shape {out.shape} and dtype {out.dtype}')
  181. half &= builder.platform_has_fast_fp16
  182. LOGGER.info(f'{prefix} building FP{16 if half else 32} engine in {f}')
  183. if half:
  184. config.set_flag(trt.BuilderFlag.FP16)
  185. with builder.build_engine(network, config) as engine, open(f, 'wb') as t:
  186. t.write(engine.serialize())
  187. LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  188. return f
  189. except Exception as e:
  190. LOGGER.info(f'\n{prefix} export failure: {e}')
  191. def export_saved_model(model, im, file, dynamic,
  192. tf_nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45,
  193. conf_thres=0.25, prefix=colorstr('TensorFlow SavedModel:')):
  194. # YOLOv5 TensorFlow SavedModel export
  195. try:
  196. import tensorflow as tf
  197. from tensorflow import keras
  198. from models.tf import TFDetect, TFModel
  199. LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
  200. f = str(file).replace('.pt', '_saved_model')
  201. batch_size, ch, *imgsz = list(im.shape) # BCHW
  202. tf_model = TFModel(cfg=model.yaml, model=model, nc=model.nc, imgsz=imgsz)
  203. im = tf.zeros((batch_size, *imgsz, 3)) # BHWC order for TensorFlow
  204. y = tf_model.predict(im, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres)
  205. inputs = keras.Input(shape=(*imgsz, 3), batch_size=None if dynamic else batch_size)
  206. outputs = tf_model.predict(inputs, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres)
  207. keras_model = keras.Model(inputs=inputs, outputs=outputs)
  208. keras_model.trainable = False
  209. keras_model.summary()
  210. keras_model.save(f, save_format='tf')
  211. LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  212. return keras_model, f
  213. except Exception as e:
  214. LOGGER.info(f'\n{prefix} export failure: {e}')
  215. return None, None
  216. def export_pb(keras_model, im, file, prefix=colorstr('TensorFlow GraphDef:')):
  217. # YOLOv5 TensorFlow GraphDef *.pb export https://github.com/leimao/Frozen_Graph_TensorFlow
  218. try:
  219. import tensorflow as tf
  220. from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
  221. LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
  222. f = file.with_suffix('.pb')
  223. m = tf.function(lambda x: keras_model(x)) # full model
  224. m = m.get_concrete_function(tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype))
  225. frozen_func = convert_variables_to_constants_v2(m)
  226. frozen_func.graph.as_graph_def()
  227. tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(f.parent), name=f.name, as_text=False)
  228. LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  229. return f
  230. except Exception as e:
  231. LOGGER.info(f'\n{prefix} export failure: {e}')
  232. def export_tflite(keras_model, im, file, int8, data, ncalib, prefix=colorstr('TensorFlow Lite:')):
  233. # YOLOv5 TensorFlow Lite export
  234. try:
  235. import tensorflow as tf
  236. LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
  237. batch_size, ch, *imgsz = list(im.shape) # BCHW
  238. f = str(file).replace('.pt', '-fp16.tflite')
  239. converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
  240. converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]
  241. converter.target_spec.supported_types = [tf.float16]
  242. converter.optimizations = [tf.lite.Optimize.DEFAULT]
  243. if int8:
  244. from models.tf import representative_dataset_gen
  245. dataset = LoadImages(check_dataset(data)['train'], img_size=imgsz, auto=False) # representative data
  246. converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib)
  247. converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
  248. converter.target_spec.supported_types = []
  249. converter.inference_input_type = tf.uint8 # or tf.int8
  250. converter.inference_output_type = tf.uint8 # or tf.int8
  251. converter.experimental_new_quantizer = False
  252. f = str(file).replace('.pt', '-int8.tflite')
  253. tflite_model = converter.convert()
  254. open(f, "wb").write(tflite_model)
  255. LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  256. return f
  257. except Exception as e:
  258. LOGGER.info(f'\n{prefix} export failure: {e}')
  259. def export_edgetpu(keras_model, im, file, prefix=colorstr('Edge TPU:')):
  260. # YOLOv5 Edge TPU export https://coral.ai/docs/edgetpu/models-intro/
  261. try:
  262. cmd = 'edgetpu_compiler --version'
  263. help_url = 'https://coral.ai/docs/edgetpu/compiler/'
  264. assert platform.system() == 'Linux', f'export only supported on Linux. See {help_url}'
  265. if subprocess.run(cmd, shell=True).returncode != 0:
  266. LOGGER.info(f'\n{prefix} export requires Edge TPU compiler. Attempting install from {help_url}')
  267. for c in ['curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -',
  268. 'echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list',
  269. 'sudo apt-get update',
  270. 'sudo apt-get install edgetpu-compiler']:
  271. subprocess.run(c, shell=True, check=True)
  272. ver = subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().split()[-1]
  273. LOGGER.info(f'\n{prefix} starting export with Edge TPU compiler {ver}...')
  274. f = str(file).replace('.pt', '-int8_edgetpu.tflite') # Edge TPU model
  275. f_tfl = str(file).replace('.pt', '-int8.tflite') # TFLite model
  276. cmd = f"edgetpu_compiler -s {f_tfl}"
  277. subprocess.run(cmd, shell=True, check=True)
  278. LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  279. return f
  280. except Exception as e:
  281. LOGGER.info(f'\n{prefix} export failure: {e}')
  282. def export_tfjs(keras_model, im, file, prefix=colorstr('TensorFlow.js:')):
  283. # YOLOv5 TensorFlow.js export
  284. try:
  285. check_requirements(('tensorflowjs',))
  286. import re
  287. import tensorflowjs as tfjs
  288. LOGGER.info(f'\n{prefix} starting export with tensorflowjs {tfjs.__version__}...')
  289. f = str(file).replace('.pt', '_web_model') # js dir
  290. f_pb = file.with_suffix('.pb') # *.pb path
  291. f_json = f + '/model.json' # *.json path
  292. cmd = f'tensorflowjs_converter --input_format=tf_frozen_model ' \
  293. f'--output_node_names="Identity,Identity_1,Identity_2,Identity_3" {f_pb} {f}'
  294. subprocess.run(cmd, shell=True)
  295. json = open(f_json).read()
  296. with open(f_json, 'w') as j: # sort JSON Identity_* in ascending order
  297. subst = re.sub(
  298. r'{"outputs": {"Identity.?.?": {"name": "Identity.?.?"}, '
  299. r'"Identity.?.?": {"name": "Identity.?.?"}, '
  300. r'"Identity.?.?": {"name": "Identity.?.?"}, '
  301. r'"Identity.?.?": {"name": "Identity.?.?"}}}',
  302. r'{"outputs": {"Identity": {"name": "Identity"}, '
  303. r'"Identity_1": {"name": "Identity_1"}, '
  304. r'"Identity_2": {"name": "Identity_2"}, '
  305. r'"Identity_3": {"name": "Identity_3"}}}',
  306. json)
  307. j.write(subst)
  308. LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  309. return f
  310. except Exception as e:
  311. LOGGER.info(f'\n{prefix} export failure: {e}')
  312. @torch.no_grad()
  313. def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
  314. weights=ROOT / 'yolov5s.pt', # weights path
  315. imgsz=(640, 640), # image (height, width)
  316. batch_size=1, # batch size
  317. device='cpu', # cuda device, i.e. 0 or 0,1,2,3 or cpu
  318. include=('torchscript', 'onnx'), # include formats
  319. half=False, # FP16 half-precision export
  320. inplace=False, # set YOLOv5 Detect() inplace=True
  321. train=False, # model.train() mode
  322. optimize=False, # TorchScript: optimize for mobile
  323. int8=False, # CoreML/TF INT8 quantization
  324. dynamic=False, # ONNX/TF: dynamic axes
  325. simplify=False, # ONNX: simplify model
  326. opset=12, # ONNX: opset version
  327. verbose=False, # TensorRT: verbose log
  328. workspace=4, # TensorRT: workspace size (GB)
  329. nms=False, # TF: add NMS to model
  330. agnostic_nms=False, # TF: add agnostic NMS to model
  331. topk_per_class=100, # TF.js NMS: topk per class to keep
  332. topk_all=100, # TF.js NMS: topk for all classes to keep
  333. iou_thres=0.45, # TF.js NMS: IoU threshold
  334. conf_thres=0.25 # TF.js NMS: confidence threshold
  335. ):
  336. t = time.time()
  337. include = [x.lower() for x in include]
  338. tf_exports = list(x in include for x in ('saved_model', 'pb', 'tflite', 'edgetpu', 'tfjs')) # TensorFlow exports
  339. file = Path(url2file(weights) if str(weights).startswith(('http:/', 'https:/')) else weights)
  340. # Checks
  341. imgsz *= 2 if len(imgsz) == 1 else 1 # expand
  342. opset = 12 if ('openvino' in include) else opset # OpenVINO requires opset <= 12
  343. # Load PyTorch model
  344. device = select_device(device)
  345. assert not (device.type == 'cpu' and half), '--half only compatible with GPU export, i.e. use --device 0'
  346. model = attempt_load(weights, map_location=device, inplace=True, fuse=True) # load FP32 model
  347. nc, names = model.nc, model.names # number of classes, class names
  348. # Input
  349. gs = int(max(model.stride)) # grid size (max stride)
  350. imgsz = [check_img_size(x, gs) for x in imgsz] # verify img_size are gs-multiples
  351. im = torch.zeros(batch_size, 3, *imgsz).to(device) # image size(1,3,320,192) BCHW iDetection
  352. # Update model
  353. if half:
  354. im, model = im.half(), model.half() # to FP16
  355. model.train() if train else model.eval() # training mode = no Detect() layer grid construction
  356. for k, m in model.named_modules():
  357. if isinstance(m, Conv): # assign export-friendly activations
  358. if isinstance(m.act, nn.SiLU):
  359. m.act = SiLU()
  360. elif isinstance(m, Detect):
  361. m.inplace = inplace
  362. m.onnx_dynamic = dynamic
  363. if hasattr(m, 'forward_export'):
  364. m.forward = m.forward_export # assign custom forward (optional)
  365. for _ in range(2):
  366. y = model(im) # dry runs
  367. LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} ({file_size(file):.1f} MB)")
  368. # Exports
  369. f = [''] * 10 # exported filenames
  370. warnings.filterwarnings(action='ignore', category=torch.jit.TracerWarning) # suppress TracerWarning
  371. if 'torchscript' in include:
  372. f[0] = export_torchscript(model, im, file, optimize)
  373. if 'engine' in include: # TensorRT required before ONNX
  374. f[1] = export_engine(model, im, file, train, half, simplify, workspace, verbose)
  375. if ('onnx' in include) or ('openvino' in include): # OpenVINO requires ONNX
  376. f[2] = export_onnx(model, im, file, opset, train, dynamic, simplify)
  377. if 'openvino' in include:
  378. f[3] = export_openvino(model, im, file)
  379. if 'coreml' in include:
  380. _, f[4] = export_coreml(model, im, file)
  381. # TensorFlow Exports
  382. if any(tf_exports):
  383. pb, tflite, edgetpu, tfjs = tf_exports[1:]
  384. if int8 or edgetpu: # TFLite --int8 bug https://github.com/ultralytics/yolov5/issues/5707
  385. check_requirements(('flatbuffers==1.12',)) # required before `import tensorflow`
  386. assert not (tflite and tfjs), 'TFLite and TF.js models must be exported separately, please pass only one type.'
  387. model, f[5] = export_saved_model(model, im, file, dynamic, tf_nms=nms or agnostic_nms or tfjs,
  388. agnostic_nms=agnostic_nms or tfjs, topk_per_class=topk_per_class,
  389. topk_all=topk_all, conf_thres=conf_thres, iou_thres=iou_thres) # keras model
  390. if pb or tfjs: # pb prerequisite to tfjs
  391. f[6] = export_pb(model, im, file)
  392. if tflite or edgetpu:
  393. f[7] = export_tflite(model, im, file, int8=int8 or edgetpu, data=data, ncalib=100)
  394. if edgetpu:
  395. f[8] = export_edgetpu(model, im, file)
  396. if tfjs:
  397. f[9] = export_tfjs(model, im, file)
  398. # Finish
  399. f = [str(x) for x in f if x] # filter out '' and None
  400. LOGGER.info(f'\nExport complete ({time.time() - t:.2f}s)'
  401. f"\nResults saved to {colorstr('bold', file.parent.resolve())}"
  402. f"\nDetect: python detect.py --weights {f[-1]}"
  403. f"\nPyTorch Hub: model = torch.hub.load('ultralytics/yolov5', 'custom', '{f[-1]}')"
  404. f"\nValidate: python val.py --weights {f[-1]}"
  405. f"\nVisualize: https://netron.app")
  406. return f # return list of exported files/dirs
  407. def parse_opt():
  408. parser = argparse.ArgumentParser()
  409. parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
  410. parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model.pt path(s)')
  411. parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640, 640], help='image (h, w)')
  412. parser.add_argument('--batch-size', type=int, default=1, help='batch size')
  413. parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
  414. parser.add_argument('--half', action='store_true', help='FP16 half-precision export')
  415. parser.add_argument('--inplace', action='store_true', help='set YOLOv5 Detect() inplace=True')
  416. parser.add_argument('--train', action='store_true', help='model.train() mode')
  417. parser.add_argument('--optimize', action='store_true', help='TorchScript: optimize for mobile')
  418. parser.add_argument('--int8', action='store_true', help='CoreML/TF INT8 quantization')
  419. parser.add_argument('--dynamic', action='store_true', help='ONNX/TF: dynamic axes')
  420. parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model')
  421. parser.add_argument('--opset', type=int, default=12, help='ONNX: opset version')
  422. parser.add_argument('--verbose', action='store_true', help='TensorRT: verbose log')
  423. parser.add_argument('--workspace', type=int, default=4, help='TensorRT: workspace size (GB)')
  424. parser.add_argument('--nms', action='store_true', help='TF: add NMS to model')
  425. parser.add_argument('--agnostic-nms', action='store_true', help='TF: add agnostic NMS to model')
  426. parser.add_argument('--topk-per-class', type=int, default=100, help='TF.js NMS: topk per class to keep')
  427. parser.add_argument('--topk-all', type=int, default=100, help='TF.js NMS: topk for all classes to keep')
  428. parser.add_argument('--iou-thres', type=float, default=0.45, help='TF.js NMS: IoU threshold')
  429. parser.add_argument('--conf-thres', type=float, default=0.25, help='TF.js NMS: confidence threshold')
  430. parser.add_argument('--include', nargs='+',
  431. default=['torchscript', 'onnx'],
  432. help='torchscript, onnx, openvino, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs')
  433. opt = parser.parse_args()
  434. print_args(FILE.stem, opt)
  435. return opt
  436. def main(opt):
  437. for opt.weights in (opt.weights if isinstance(opt.weights, list) else [opt.weights]):
  438. run(**vars(opt))
  439. if __name__ == "__main__":
  440. opt = parse_opt()
  441. main(opt)