TensorRT Export Fix (#171)

README.md

@@ -18,11 +18,9 @@
   </div>
   <br>
 
-[Ultralytics YOLOv8](https://github.com/ultralytics/ultralytics) is the latest version of the YOLO object detection and image segmentation model developed by [Ultralytics](https://ultralytics.com). YOLOv8 is a cutting-edge, state-of-the-art (SOTA) model that builds upon the success of previous YOLO versions and introduces new features and improvements to further boost performance and flexibility.
+[Ultralytics YOLOv8](https://github.com/ultralytics/ultralytics), developed by [Ultralytics](https://ultralytics.com), is a cutting-edge, state-of-the-art (SOTA) model that builds upon the success of previous YOLO versions and introduces new features and improvements to further boost performance and flexibility. YOLOv8 is designed to be fast, accurate, and easy to use, making it an excellent choice for a wide range of object detection, image segmentation and image classification tasks.
 
-The YOLOv8 models are designed to be fast, accurate, and easy to use, making them an excellent choice for a wide range of object detection, image segmentation and image classification tasks.
+To request an Enterprise License please complete the form at [Ultralytics Licensing](https://ultralytics.com/license).
-
-Whether you are a seasoned machine learning practitioner or new to the field, we hope that the resources on this page will help you get the most out of YOLOv8.
 
 <div align="center">
     <a href="https://github.com/ultralytics" style="text-decoration:none;">
@@ -101,8 +99,9 @@ Ultralytics [release](https://github.com/ultralytics/ultralytics/releases).
 
 ### Known Issues / TODOs
 
+We are still working on several parts of YOLOv8! We aim to have these completed soon to bring the YOLOv8 feature set up to par with YOLOv5, including export and inference to all the same formats. We are also writing a YOLOv8 paper which we will submit to [arxiv.org](https://arxiv.org) once complete.
+
 - [ ] TensorFlow exports
-- [ ] GPU exports
 - [ ] DDP resume
 - [ ] [arxiv.org](https://arxiv.org) paper
 
@@ -170,7 +169,7 @@ Ultralytics [release](https://github.com/ultralytics/ultralytics/releases) on fi
 
 <br>
 <a align="center" href="https://bit.ly/ultralytics_hub" target="_blank">
-<img width="100%" src="https://github.com/ultralytics/assets/raw/main/im/integrations-loop.png"></a>
+<img width="100%" src="https://github.com/ultralytics/assets/raw/main/yolov8/banner-integrations.png"></a>
 <br>
 <br>
 

examples/tutorial.ipynb

@@ -290,7 +290,7 @@
       "source": [
         "# 3. Train\n",
         "\n",
-        "<p align=\"\"><a href=\"https://roboflow.com/?ref=ultralytics\"><img width=\"1000\" src=\"https://github.com/ultralytics/assets/raw/main/im/integrations-loop.png\"/></a></p>\n",
+        "<p align=\"\"><a href=\"https://roboflow.com/?ref=ultralytics\"><img width=\"1000\" src=\"https://github.com/ultralytics/assets/raw/main/yolov8/banner-integrations.png\"/></a></p>\n",
         "\n",
         "Train YOLOv8 on detection, segmentation and classification datasets."
       ]

ultralytics/yolo/configs/default.yaml

@@ -13,7 +13,7 @@ batch: 16 # number of images per batch
 imgsz: 640 # size of input images
 save: True # save checkpoints
 cache: False # True/ram, disk or False. Use cache for data loading
-device: '' # cuda device, i.e. 0 or 0,1,2,3 or cpu. Device to run on
+device: null # cuda device, i.e. 0 or 0,1,2,3 or cpu. Device to run on
 workers: 8 # number of worker threads for data loading
 project: null # project name
 name: null # experiment name

ultralytics/yolo/engine/exporter.py

@@ -152,7 +152,7 @@ class Exporter:
         jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle = flags  # export booleans
 
         # Load PyTorch model
-        self.device = select_device(self.args.device or 'cpu')
+        self.device = select_device('cpu' if self.args.device is None else self.args.device)
         if self.args.half:
             if self.device.type == 'cpu' and not coreml:
                 LOGGER.info('half=True only compatible with GPU or CoreML export, i.e. use device=0 or format=coreml')
@@ -173,7 +173,7 @@ class Exporter:
             file = Path(file.name)
 
         # Update model
-        model = deepcopy(model)
+        model = deepcopy(model).to(self.device)
         for p in model.parameters():
             p.requires_grad = False
         model.eval()
@@ -218,6 +218,8 @@ class Exporter:
         if coreml:  # CoreML
             f[4], _ = self._export_coreml()
         if any((saved_model, pb, tflite, edgetpu, tfjs)):  # TensorFlow formats
+            raise NotImplementedError('YOLOv8 TensorFlow export support is still under development. '
+                                      'Please consider contributing to the effort if you have TF expertise. Thank you!')
             assert not isinstance(model, ClassificationModel), 'ClassificationModel TF exports not yet supported.'
             nms = False
             f[5], s_model = self._export_saved_model(nms=nms or self.args.agnostic_nms or tfjs,