Fix Python3.9 CI and reduce unit test memory cost of PVT & Swin (#6366)

* Fix py3.9 CI.

* Fix py3.9 CI.

* change ubuntu version

* use install python

* add 3.9 source

* remove sudo

* fix

* fix

* noninteractive

* fix env

* fix all ci

* add comments

* reduce swin and pvt input size

* resolve comments

.github/workflows/build.yml

@@ -167,7 +167,7 @@ jobs:
 
     strategy:
       matrix:
-        python-version: [3.6, 3.7, 3.8, 3.9-dev]
+        python-version: [3.6, 3.7, 3.8, 3.9]
         torch: [1.9.0+cu102]
         include:
           - torch: 1.9.0+cu102
@@ -181,9 +181,17 @@ jobs:
         uses: actions/setup-python@v2
         with:
           python-version: ${{ matrix.python-version }}
+      # Add ppa source repo for python3.9.
+      - name: Add python3.9 source
+        run: |
+          apt-get update && apt-get install -y software-properties-common
+          add-apt-repository -y ppa:deadsnakes/ppa
+        if: ${{matrix.python-version == '3.9'}}
+      # Install python-dev for some packages which require libpython3.Xm.
+      # Github's setup-python cannot install python3.9-dev, so we have to use apt install.
+      # Set DEBIAN_FRONTEND=noninteractive to avoid some interactions.
       - name: Install python-dev
-        run: apt-get update && apt-get install -y python${{matrix.python-version}}-dev
-        if: ${{matrix.python-version != '3.9-dev'}}
+        run: apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends python${{matrix.python-version}}-dev
       - name: Install system dependencies
         run: |
           apt-get update && apt-get install -y ffmpeg libsm6 libxext6 git ninja-build libglib2.0-0 libsm6 libxrender-dev libxext6
@@ -196,7 +204,7 @@ jobs:
         run: python -m pip install torch==${{matrix.torch}} torchvision==${{matrix.torchvision}} -f https://download.pytorch.org/whl/torch_stable.html
       - name: Install dependencies for compiling onnx when python=3.9
         run: python -m pip install protobuf && apt-get update && apt-get -y install libprotobuf-dev protobuf-compiler cmake
-        if: ${{matrix.python-version == '3.9-dev'}}
+        if: ${{matrix.python-version == '3.9'}}
       - name: Install mmdet dependencies
         run: |
           python -V

tests/test_models/test_backbones/test_pvt.py

@@ -37,22 +37,22 @@ def test_pvt():
     model(temp)
 
     # Test normal inference
-    temp = torch.randn((1, 3, 512, 512))
+    temp = torch.randn((1, 3, 32, 32))
     model = PyramidVisionTransformer()
     outs = model(temp)
-    assert outs[0].shape == (1, 64, 128, 128)
-    assert outs[1].shape == (1, 128, 64, 64)
-    assert outs[2].shape == (1, 320, 32, 32)
-    assert outs[3].shape == (1, 512, 16, 16)
+    assert outs[0].shape == (1, 64, 8, 8)
+    assert outs[1].shape == (1, 128, 4, 4)
+    assert outs[2].shape == (1, 320, 2, 2)
+    assert outs[3].shape == (1, 512, 1, 1)
 
     # Test abnormal inference size
-    temp = torch.randn((1, 3, 511, 511))
+    temp = torch.randn((1, 3, 33, 33))
     model = PyramidVisionTransformer()
     outs = model(temp)
-    assert outs[0].shape == (1, 64, 127, 127)
-    assert outs[1].shape == (1, 128, 63, 63)
-    assert outs[2].shape == (1, 320, 31, 31)
-    assert outs[3].shape == (1, 512, 15, 15)
+    assert outs[0].shape == (1, 64, 8, 8)
+    assert outs[1].shape == (1, 128, 4, 4)
+    assert outs[2].shape == (1, 320, 2, 2)
+    assert outs[3].shape == (1, 512, 1, 1)
 
     # Test abnormal inference size
     temp = torch.randn((1, 3, 112, 137))
@@ -76,22 +76,22 @@ def test_pvtv2():
         PyramidVisionTransformerV2(pretrain_img_size=(224, 224, 224))
 
     # Test normal inference
-    temp = torch.randn((1, 3, 512, 512))
+    temp = torch.randn((1, 3, 32, 32))
     model = PyramidVisionTransformerV2()
     outs = model(temp)
-    assert outs[0].shape == (1, 64, 128, 128)
-    assert outs[1].shape == (1, 128, 64, 64)
-    assert outs[2].shape == (1, 320, 32, 32)
-    assert outs[3].shape == (1, 512, 16, 16)
+    assert outs[0].shape == (1, 64, 8, 8)
+    assert outs[1].shape == (1, 128, 4, 4)
+    assert outs[2].shape == (1, 320, 2, 2)
+    assert outs[3].shape == (1, 512, 1, 1)
 
     # Test abnormal inference size
-    temp = torch.randn((1, 3, 511, 511))
+    temp = torch.randn((1, 3, 31, 31))
     model = PyramidVisionTransformerV2()
     outs = model(temp)
-    assert outs[0].shape == (1, 64, 128, 128)
-    assert outs[1].shape == (1, 128, 64, 64)
-    assert outs[2].shape == (1, 320, 32, 32)
-    assert outs[3].shape == (1, 512, 16, 16)
+    assert outs[0].shape == (1, 64, 8, 8)
+    assert outs[1].shape == (1, 128, 4, 4)
+    assert outs[2].shape == (1, 320, 2, 2)
+    assert outs[3].shape == (1, 512, 1, 1)
 
     # Test abnormal inference size
     temp = torch.randn((1, 3, 112, 137))

tests/test_models/test_backbones/test_swin.py

@@ -50,22 +50,22 @@ def test_swin_transformer():
     model(temp)
 
     # Test normal inference
-    temp = torch.randn((1, 3, 512, 512))
+    temp = torch.randn((1, 3, 32, 32))
     model = SwinTransformer()
     outs = model(temp)
-    assert outs[0].shape == (1, 96, 128, 128)
-    assert outs[1].shape == (1, 192, 64, 64)
-    assert outs[2].shape == (1, 384, 32, 32)
-    assert outs[3].shape == (1, 768, 16, 16)
+    assert outs[0].shape == (1, 96, 8, 8)
+    assert outs[1].shape == (1, 192, 4, 4)
+    assert outs[2].shape == (1, 384, 2, 2)
+    assert outs[3].shape == (1, 768, 1, 1)
 
     # Test abnormal inference size
-    temp = torch.randn((1, 3, 511, 511))
+    temp = torch.randn((1, 3, 31, 31))
     model = SwinTransformer()
     outs = model(temp)
-    assert outs[0].shape == (1, 96, 128, 128)
-    assert outs[1].shape == (1, 192, 64, 64)
-    assert outs[2].shape == (1, 384, 32, 32)
-    assert outs[3].shape == (1, 768, 16, 16)
+    assert outs[0].shape == (1, 96, 8, 8)
+    assert outs[1].shape == (1, 192, 4, 4)
+    assert outs[2].shape == (1, 384, 2, 2)
+    assert outs[3].shape == (1, 768, 1, 1)
 
     # Test abnormal inference size
     temp = torch.randn((1, 3, 112, 137))