[cherry-pick] fixes from release/0.2 testing

README.md

@@ -312,7 +312,7 @@ clone the repository and install any additional dependencies as required.
 
 The currently supported framework versions are:
 
-- PyTorch supported versions: `>= 1.1.0, < 1.7.0`
+- PyTorch supported versions: `>= 1.1.0, < 1.8.0`
 - Keras supported versions: `2.3.0-tf` (through the TensorFlow `2.2` package; as of Feb 1st, 2021, `keras2onnx` has
 not been tested for TensorFlow >= `2.3`). 
 - TensorFlow V1 supported versions: >= `1.8.0` (TensorFlow >= `2.X` is not currently supported)

integrations/ultralytics/README.md

@@ -44,8 +44,8 @@ cp ../sparseml/integrations/ultralytics/*.py .
 cp ../sparseml/integrations/ultralytics/deepsparse/*.py .
 
 # install dependencies
+pip install sparseml[torchvision] deepsparse
 pip install -r requirements.txt
-pip install sparseml deepsparse
 ```
 
 
@@ -119,7 +119,7 @@ cp sparseml/integrations/ultralytics/deepsparse/*.py yolov5
 cd yolov5
 
 # install deepsparse and server dependencies
-pip install deepsparse sparseml flask flask-cors
+pip install deepsparse sparseml[torchvision] flask flask-cors
 ```
 
 Note: on new Ubuntu systems, to install `cv2` running `sudo apt-get update && apt-get install -y python3-opencv`

integrations/ultralytics/deepsparse/SERVER.md

@@ -41,8 +41,8 @@ cp sparseml/integrations/ultralytics/deepsparse/*.py yolov5
 cd yolov5
 
 # install dependencies
+pip install deepsparse sparseml[torchvision] flask flask-cors
 pip install -r requirements.txt
-pip install deepsparse sparseml flask flask-cors
 ```
 
 ## Execution

integrations/ultralytics/deepsparse/deepsparse_utils.py

@@ -74,7 +74,7 @@ class YoloPostprocessor:
         output shapes
     """
 
-    def __init__(self, image_size: Tuple[int]):
+    def __init__(self, image_size: Tuple[int] = (640, 640)):
         self._image_size = image_size
         self._grids = {}  # Dict[Tuple[int], torch.Tensor]
 

integrations/ultralytics/deepsparse/server.py

@@ -56,7 +56,7 @@
 import flask
 from deepsparse import compile_model
 from deepsparse.utils import arrays_to_bytes, bytes_to_arrays
-from deepsparse_utils import postprocess_nms, pre_nms_postprocess
+from deepsparse_utils import YoloPostprocessor, postprocess_nms
 from flask_cors import CORS
 
 
@@ -121,6 +121,8 @@ def create_and_run_model_server(
     engine = compile_model(model_path, batch_size, num_cores)
     print(engine)
 
+    postprocessor = YoloPostprocessor()
+
     app = flask.Flask(__name__)
     CORS(app)
 
@@ -145,7 +147,7 @@ def predict():
 
         # post-processing
         postprocess_start_time = time.time()
-        outputs = pre_nms_postprocess(outputs)
+        outputs = postprocessor.pre_nms_postprocess(outputs)
         postprocess_time = time.time() - postprocess_start_time
         print(f"Post-processing, pre-nms time: {postprocess_time * 1000.0:.4f}ms")
 

setup.py

@@ -55,8 +55,8 @@
 _deepsparse_deps = [
     f"{'deepsparse-nightly' if _NIGHTLY else 'deepsparse'}~={_VERSION_MAJOR_MINOR}"
 ]
-_pytorch_deps = ["torch>=1.1.0", "tensorboard>=1.0", "tensorboardX>=1.0"]
-_pytorch_vision_deps = _pytorch_deps + ["torchvision>=0.3.0"]
+_pytorch_deps = ["torch>=1.1.0,<1.8", "tensorboard>=1.0", "tensorboardX>=1.0"]
+_pytorch_vision_deps = _pytorch_deps + ["torchvision>=0.3.0,<0.9"]
 _tensorflow_v1_deps = ["tensorflow<2.0.0", "tensorboard<2.0.0", "tf2onnx>=1.0.0,<1.6"]
 _tensorflow_v1_gpu_deps = [
     "tensorflow-gpu<2.0.0",