Fixes ORT segfault in nightly version (#18)

* Fixes onnx runtime segfault

* Fix libtorch header

* Update with torchvision interface

* Abandon C++ test of stable version

.github/workflows/nightly.yml

@@ -33,6 +33,7 @@ jobs:
     - name: Install PyTorch Nightly
       run: |
         pip install --pre torch torchvision -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
+        pip install --user -i https://test.pypi.org/simple/ ort-nightly==1.5.2.dev202012031
     - name: Test with pytest
       run: |
         pytest test

.github/workflows/stable.yml

@@ -30,34 +30,10 @@ jobs:
         python -m pip install --upgrade pip
         pip install flake8 pytest
         if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
-    - name: Install PyTorch 1.7
+    - name: Install PyTorch 1.7.1
       run: |
-        pip install torch==1.7.0+cpu torchvision==0.8.1+cpu -f https://download.pytorch.org/whl/torch_stable.html
+        pip install torch==1.7.1+cpu torchvision==0.8.2+cpu -f https://download.pytorch.org/whl/torch_stable.html
+        pip install onnxruntime
     - name: Test with pytest
       run: |
         pytest test
-    - name: Build TorchVision Cpp 0.8.0
-      run: |
-        export TORCH_PATH=$(dirname $(python -c "import torch; print(torch.__file__)"))
-        cd ..
-        git clone https://github.com/pytorch/vision.git vision
-        cd vision
-        git checkout release/0.8.0
-        mkdir build && cd build
-        cmake .. -DTorch_DIR=$TORCH_PATH/share/cmake/Torch
-        make -j4
-        sudo make install
-    - name: Export torchscript model
-      run: |
-        python -m test.tracing.trace_model
-    - name: Test libtorch tracing
-      run: |
-        export TORCH_PATH=$(dirname $(python -c "import torch; print(torch.__file__)"))
-        export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$TORCH_PATH/lib/
-        cd test/tracing
-        mkdir build && cd build
-        cmake .. -DTorch_DIR=$TORCH_PATH/share/cmake/Torch
-        make
-        mv ../yolov5s.torchscript.pt ./
-        echo ">> Test libtorch tracing"
-        ./test_tracing

notebooks/export-onnx-inference-onnxruntime.ipynb

@@ -242,7 +242,7 @@
     "    export_onnx_name,\n",
     "    do_constant_folding=True,\n",
     "    opset_version=_onnx_opset_version,\n",
-    "    dynamic_axes={\"images_tensors\": [0, 1, 2, 3], \"outputs\": [0, 1, 2, 3]}, \n",
+    "    dynamic_axes={\"images_tensors\": [0, 1, 2], \"outputs\": [0, 1, 2]}, \n",
     "    input_names=[\"images_tensors\"],\n",
     "    output_names=[\"outputs\"],\n",
     ")"

requirements.txt

@@ -10,6 +10,6 @@ scipy>=1.4.1
 tensorboard>=2.2
 tqdm>=4.41.0
 onnx>=1.8.0
-onnxruntime>=1.5.2
+# onnxruntime>=1.5.2  # pytorch nightly needs to rely on matching with ort-nightly
 # pycocotools need python3.7 as minimal
 # pip install -U pycocotools>=2.0.2  # corresponds to https://github.com/ppwwyyxx/cocoapi

test/test_onnx.py

@@ -102,12 +102,12 @@ def test_yolov5s(self):
         # Test exported model on images of different size, or dummy input
         self.run_model(model, [(images,), (test_images,), (dummy_image,)], input_names=["images_tensors"],
                        output_names=["outputs"],
-                       dynamic_axes={"images_tensors": [0, 1, 2, 3], "outputs": [0, 1, 2, 3]},
+                       dynamic_axes={"images_tensors": [0, 1, 2], "outputs": [0, 1, 2]},
                        tolerate_small_mismatch=True)
         # Test exported model for an image with no detections on other images
         self.run_model(model, [(dummy_image,), (images,)], input_names=["images_tensors"],
                        output_names=["outputs"],
-                       dynamic_axes={"images_tensors": [0, 1, 2, 3], "outputs": [0, 1, 2, 3]},
+                       dynamic_axes={"images_tensors": [0, 1, 2], "outputs": [0, 1, 2]},
                        tolerate_small_mismatch=True)
 
 

test/tracing/test_tracing.cpp

@@ -1,9 +1,7 @@
-#include <ATen/ATen.h>
 #include <torch/script.h>
 #include <torch/torch.h>
 
 #include <torchvision/vision.h>
-#include <torchvision/ROIPool.h>
 #include <torchvision/nms.h>