Add PyTorch 1.10 to GH Actions (#207)

* Resolve torchvision version conflict

* Add PyTorch 1.10 to GH Actions

* Minor fixes

* Resolve torchvision version conflict

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

.github/workflows/ci_test.yml

@@ -11,16 +11,21 @@ on:
 
 jobs:
   UnitTest:
-    runs-on: ${{ matrix.os }}
+    runs-on: ${{ matrix.image }}
     strategy:
       matrix:
-        python-version: [3.7]
-        os: [ubuntu-latest]
-        torch: [1.9.1+cpu]
+        python-version: [ 3.7 ]
+        image: [ 'ubuntu-latest' ]
+        torch: [ '1.9.1+cpu', '1.10.0+cpu' ]
         include:
-          - torch: 1.9.1+cpu
-            pip_address: torch==1.9.1+cpu torchvision==0.10.1+cpu -f https://download.pytorch.org/whl/torch_stable.html
+          - torch: '1.9.1+cpu'
+            torch_address: torch==1.9.1+cpu torchvision==0.10.1+cpu -f https://download.pytorch.org/whl/torch_stable.html
             unittest_type: --cov=test --cov-report=xml
+            torchvision: release/0.10
+          - torch: '1.10.0+cpu'
+            torch_address: install torch==1.10.0+cpu torchvision==0.11.1+cpu -f https://download.pytorch.org/whl/cpu/torch_stable.html
+            unittest_type: --cov=test --cov-report=xml
+            torchvision: release/0.11
 
     steps:
       - name: Clone repository
@@ -35,7 +40,7 @@ jobs:
           python -m pip install --upgrade pip
           # requirements for PyTorch and torchvision
           pip install numpy pillow scipy
-          pip install ${{ matrix.pip_address }}
+          pip install ${{ matrix.torch_address }}
       - name: Install dependencies
         run: |
           # requirements for unittest
@@ -71,7 +76,7 @@ jobs:
           cd ..
           git clone https://github.com/pytorch/vision.git vision
           cd vision
-          git checkout release/0.10
+          git checkout ${{ matrix.torchvision }}
           mkdir build && cd build
           cmake .. -DTorch_DIR=$TORCH_PATH/share/cmake/Torch
           make -j4

yolort/data/transforms.py

@@ -10,6 +10,14 @@
 from torchvision.transforms import functional as F
 from torchvision.transforms import transforms as T
 
+try:
+    from torchvision.transforms.functional import get_image_size, get_image_num_channels
+except ImportError:
+    from torchvision.transforms.functional import (
+        _get_image_size as get_image_size,
+        _get_image_num_channels as get_image_num_channels,
+    )
+
 
 def collate_fn(batch):
     return tuple(zip(*batch))
@@ -51,7 +59,7 @@ def forward(
         if torch.rand(1) < self.p:
             image = F.hflip(image)
             if target is not None:
-                width, _ = F._get_image_size(image)
+                width, _ = get_image_size(image)
                 target["boxes"][:, [0, 2]] = width - target["boxes"][:, [2, 0]]
                 if "masks" in target:
                     target["masks"] = target["masks"].flip(-1)
@@ -106,7 +114,7 @@ def forward(
             elif image.ndimension() == 2:
                 image = image.unsqueeze(0)
 
-        orig_w, orig_h = F._get_image_size(image)
+        orig_w, orig_h = get_image_size(image)
 
         while True:
             # sample an option
@@ -206,7 +214,7 @@ def forward(
         if torch.rand(1) < self.p:
             return image, target
 
-        orig_w, orig_h = F._get_image_size(image)
+        orig_w, orig_h = get_image_size(image)
 
         r = self.side_range[0] + torch.rand(1) * (
             self.side_range[1] - self.side_range[0]
@@ -291,7 +299,7 @@ def forward(
                 image = self._contrast(image)
 
         if r[6] < self.p:
-            channels = F._get_image_num_channels(image)
+            channels = get_image_num_channels(image)
             permutation = torch.randperm(channels)
 
             is_pil = F._is_pil_image(image)

yolort/models/darknetv5.py

@@ -3,8 +3,8 @@
 
 import torch
 from torch import nn, Tensor
-from torch.hub import load_state_dict_from_url
 
+from yolort.utils import load_state_dict_from_url
 from yolort.v5 import Conv, Focus, BottleneckCSP, C3, SPPF
 from ._utils import _make_divisible
 

yolort/models/darknetv6.py

@@ -3,8 +3,8 @@
 
 import torch
 from torch import nn, Tensor
-from torch.hub import load_state_dict_from_url
 
+from yolort.utils import load_state_dict_from_url
 from yolort.v5 import Conv, C3
 from ._utils import _make_divisible
 

yolort/models/yolo.py

@@ -5,9 +5,8 @@
 
 import torch
 from torch import nn, Tensor
-from torchvision.models.utils import load_state_dict_from_url
 
-from yolort.utils import load_from_ultralytics
+from yolort.utils import load_from_ultralytics, load_state_dict_from_url
 from .anchor_utils import AnchorGenerator
 from .backbone_utils import darknet_pan_backbone
 from .box_head import YOLOHead, SetCriterion, PostProcess

yolort/utils/__init__.py

@@ -1,5 +1,10 @@
 from typing import Callable, Dict, Mapping, Sequence, Union
 
+try:
+    from torch.hub import load_state_dict_from_url
+except ImportError:
+    from torch.utils.model_zoo import load_url as load_state_dict_from_url
+
 from .hooks import FeatureExtractor
 from .image_utils import cv2_imshow, get_image_from_url, read_image_to_tensor
 from .update_module_state import load_from_ultralytics
@@ -11,6 +16,7 @@
     "get_image_from_url",
     "get_callable_dict",
     "load_from_ultralytics",
+    "load_state_dict_from_url",
     "read_image_to_tensor",
 ]