Add TensorFlow formats to export.py

[glenn-jocher]

@zldrobit this my main PR for migrating TensorFlow export support to export.py following #1127, while retaining the Keras model-building functionality you created in tf.py.

Export a YOLOv5 PyTorch model to TorchScript, ONNX, CoreML, TensorFlow (saved_model, pb, TFLite, TF.js,) formats
TensorFlow exports authored by https://github.com/zldrobit

Usage:

$ python path/to/export.py --weights yolov5s.pt --include torchscript onnx coreml saved_model pb tflite tfjs

Inference:

$ python path/to/detect.py --weights yolov5s.pt
                                     yolov5s.onnx  (must export with --dynamic)
                                     yolov5s_saved_model
                                     yolov5s.pb
                                     yolov5s.tflite

TensorFlow.js:

$ python path/to/export.py --weights yolov5s.pt --include tfjs
$ # Edit yolov5s_web_model/model.json to sort Identity* in ascending order
$ cd .. && git clone https://github.com/zldrobit/tfjs-yolov5-example.git && cd tfjs-yolov5-example
$ npm install
$ ln -s ../../yolov5/yolov5s_web_model public/yolov5s_web_model
$ npm start

🛠️ PR Summary

_{Made with ❤️ by Ultralytics Actions}

🌟 Summary

This PR updates the Ultralytics YOLOv5 repository with new export capabilities and various adjustments.

📊 Key Changes

• Added *_web_model/ to .dockerignore and .gitignore, suggesting new support for web models.
• Updated dependencies for export functionality, replacing onnx-simplifier and coremltools with tensorflow-cpu.
• Introduced changes in CI testing to accommodate new export pipelines.
• Modified detect.py and export.py for stronger TensorFlow integration and simplifications within the export process.
• Improved the tf.py TensorFlow model script with new functionalities for exporting models, removing redundant code, and introducing TF.js export support.
• Added new models (e.g., TFBN, TFConv, etc.) that seem to correspond to TensorFlow/Keras equivalents of existing YOLOv5 layers.
• Simplified the requirements in requirements.txt, highlighting the shift towards TensorFlow-based operations.

🎯 Purpose & Impact

• 🚀 Enhanced Export Options: Users can now export their YOLOv5 models to a variety of TensorFlow formats, facilitating integration into web applications and providing a better pathway for model deployment on various platforms.
• 🛠 Improved Codebase: By adding support for TensorFlow.js and refactoring the TensorFlow converting script, the codebase becomes more versatile and user-friendly for those familiar with the TF ecosystem.
• ☁️ Easier Cloud Integration: The changes in .github/workflows/ci-testing.yml imply that CI now includes tests for TensorFlow export functionality, which ensures reliability for cloud-based applications.
• ⚡ Streamlined Workflows: Removing dependencies on onnx-simplifier and coremltools may lower the barrier to entry for some users and streamline the model export process, focusing more on TensorFlow-centric workflows.

[glenn-jocher]

@zldrobit this PR is almost done! I'm understanding a lot more about TF export after doing all of this, and I had a few quick questions:

• Fusing YOLOv5 models before TF export. I can't get this to work at all with TF exports. The rest of the exports work well with or without fusing (i.e. ONNX and CoreML models export correctly either way). TF fused export does not throw an error, but detect.py does not detect anything with any fused TF models (pb, tflite, saved_model). I thought I had accounted for fused Conv() layers in this line here. Do you know what the problem might be? It would be advantageous (for faster inference and smaller size) to be able to export fused TF YOLOv5 models.
  

  yolov5/models/tf.py

  Line 99 in cd810c8

  self.bn = tf_BN(w.bn) if hasattr(w, 'bn') else tf.identity

• Trainable Params. I think we want zero trainable parameters during the export process. I assume these have gradients enabled. Do you know how to disable these in TF before export?

==================================================================================================
Total params: 7,295,869
Trainable params: 7,276,605
Non-trainable params: 19,264
__________________________________________________________________________________________________

• untraced functions. During saved_model export I get this warning, not sure if it's important or not:

Found untraced functions such as tf__conv_layer_call_fn, tf__conv_layer_call_and_return_conditional_losses, tf_bn_1_layer_call_fn, tf_bn_1_layer_call_and_return_conditional_losses, tf__conv_2_layer_call_fn while saving (showing 5 of 845). These functions will not be directly callable after loading.

• yolov5s-fp16.tflite seems to be in FP8 or INT8. When I look at the filesizes of yolov5s-fp16.tflite and yolov5s-int8.tflite they are both about 7.7 MB, about 50% the size of yolov5s.pt, which is saved in FP16 by default. So it seems both TFLite models are already quantized to 8 bits, and the 'fp16' label is incorrect?

[zldrobit]

@glenn-jocher I did some tests to verify the solutions for the following questions:

• Fusing YOLOv5 models before TF export. This is because in the construction of conv,
  bias_initilizer is set to zeros by default (https://www.tensorflow.org/api_docs/python/tf/keras/layers/Conv2D). To successfully convert a fused model,
  

  yolov5/models/tf.py

  Lines 95 to 97 in cd810c8

  	conv = keras.layers.Conv2D(
  	c2, k, s, 'SAME' if s == 1 else 'VALID', use_bias=False,
  	kernel_initializer=keras.initializers.Constant(w.conv.weight.permute(2, 3, 1, 0).numpy()))

  
  has to be changed to

conv = keras.layers.Conv2D(                                                                            
    c2, k, s, 'SAME' if s == 1 else 'VALID', use_bias=False if hasattr(w, 'bn') else True,             
    kernel_initializer=keras.initializers.Constant(w.conv.weight.permute(2, 3, 1, 0).numpy()),         
    bias_initializer='zeros' if hasattr(w, 'bn') else keras.initializers.Constant(w.conv.bias.numpy()))

• Trainable Params. Set model.trainable=False could turn all parms to non-trainable params (https://www.tensorflow.org/guide/keras/transfer_learning#recursive_setting_of_the_trainable_attribute), as in export.py

         keras_model = keras.Model(inputs=inputs, outputs=outputs)
         keras_model.trainable = False                            
         keras_model.summary()                                    
         keras_model.save(f, save_format='tf')                    

• untraced functions. This warning message can be ignored, as discussed in https://groups.google.com/a/tensorflow.org/g/tflite/c/WWrgz3Nh1TE
  The above warning message can be ignored. The warning message says that maybe there are non callable functions stored in the graph. quoted from the above link.

• yolov5s-fp16.tflite To specify fp16 precision in TFLite conversion, we have to add

converter.target_spec.supported_types = [tf.float16]

in export_tflite() of export.py. Note that the inference time per image on an i9 CPU reduced significantly from ~4s (320x320 input) to less than 400ms (640x640 input) after adding this line.

For convenience, I'll send an PR to address these issues.

[glenn-jocher]

@zldrobit wow thanks for the investigation! That's awesome you found a fuse solution. I'll take a deeper look later but to give you a quick update, tflite export (and tfjs exports seem to be working well now via export.py, usage is like this:)

python export.py --weights yolov5s.pt --include tflite  # produces yolov5s.tflite which was previously labelled as fp16
python export.py --weights yolov5s.pt --include tflite --int8  # produces yolov5s-int8.tflite using dataset generator
python export.py --weights yolov5s.pt --include tfjs  # produces yolov5s_web_model

EDIT: BTW I was not able to pass all parameters to all places in the TF model, agnostic NMS for example now uses hard coded NMS parameters, but I plan to try to fix this in the future. Previously these parameters were global variables from opt, now they need to be passed in as local variables.

[zldrobit]

@glenn-jocher I found a solution to pass all parameters to the TF NMS operation in #4905. That PR also addresses the Identity_* reordering problem by regex substitution when exporting tfjs models, so no manual operation is needed.

[glenn-jocher]

@zldrobit nice work! I've merged #4905 now.