Refactor test code

Author: atalman

.github/workflows/validate-linux-binaries.yml

@@ -29,8 +29,6 @@ jobs:
       package-type: libtorch
       os: linux
       channel: nightly
-
-
   validate-linux-binaries-conda:
     needs: generate-conda-matrix
     strategy:
@@ -39,11 +37,8 @@ jobs:
       fail-fast: false
     runs-on: ${{ matrix.validation_runner }}
     steps:
-      - name: Checkout PyTorch atalman/builder
+      - name: Checkout PyTorch builder
         uses: actions/checkout@v2
-        with:
-          ref: main
-          repository: atalman/builder
       - name: Validate binary conda
         uses: ./.github/actions/validate-linux-binary
         with:

test/smoke_test/smoke_test.py

@@ -3,27 +3,25 @@
 import torch
 import torchvision
 import torchaudio
+import pathlib
 
 gpu_arch_ver = os.getenv('GPU_ARCH_VER')
 gpu_arch_type = os.getenv('GPU_ARCH_TYPE')
-is_cuda_system = gpu_arch_type == "cuda"
+is_cuda_system = gpu_arch_type == 'cuda'
 
 def smoke_test_cuda() -> None:
     if(not torch.cuda.is_available() and is_cuda_system):
-        print(f"Expected CUDA {gpu_arch_ver}. However CUDA is not loaded.")
-        sys.exit(1)
+        raise RuntimeError(f'Expected CUDA {gpu_arch_ver}. However CUDA is not loaded.')
     if(torch.cuda.is_available()):
         if(torch.version.cuda != gpu_arch_ver):
-            print(f"Wrong CUDA version. Loaded: {torch.version.cuda} Expected: {gpu_arch_ver}")
-            sys.exit(1)
-        y=torch.randn([3,5]).cuda()
-        print(f"torch cuda: {torch.version.cuda}")
+            raise RuntimeError(f'Wrong CUDA version. Loaded: {torch.version.cuda} Expected: {gpu_arch_ver}')
+        print(f'torch cuda: {torch.version.cuda}')
         #todo add cudnn version validation
-        print(f"torch cudnn: {torch.backends.cudnn.version()}")
+        print(f'torch cudnn: {torch.backends.cudnn.version()}')
 
 def smoke_test_conv2d() -> None:
     import torch.nn as nn
-    print(f"Calling smoke_test_conv2d")
+    print('Calling smoke_test_conv2d')
     # With square kernels and equal stride
     m = nn.Conv2d(16, 33, 3, stride=2)
     # non-square kernels and unequal stride and with padding
@@ -33,29 +31,28 @@ def smoke_test_conv2d() -> None:
     input = torch.randn(20, 16, 50, 100)
     output = m(input)
     if(is_cuda_system):
-        print(f"Testing smoke_test_conv2d with cuda")
+        print('Testing smoke_test_conv2d with cuda')
         conv = nn.Conv2d(3, 3, 3).cuda()
         x = torch.randn(1, 3, 24, 24).cuda()
         with torch.cuda.amp.autocast():
             out = conv(x)
-        print(out.sum())
 
 def smoke_test_torchvision() -> None:
-    import torchvision.datasets as dset
-    import torchvision.transforms
-    from torchvision.io import read_file, decode_jpeg, decode_png
     print('Is torchvision useable?', all(x is not None for x in [torch.ops.image.decode_png, torch.ops.torchvision.roi_align]))
-    img_jpg = read_file('./test/smoke_test/assets/rgb_pytorch.jpg')
-    img_jpg_nv = decode_jpeg(img_jpg)
-    img_png = read_file('./test/smoke_test/assets/rgb_pytorch.png')
-    img__png_nv = decode_png(img_png)
 
-def smoke_test_vision_resnet50_classify() -> None:
+def smoke_test_torchvision_read_decode() -> None:
+    from torchvision.io import decode_image, read_image,
+    img_jpg = read_image(str(pathlib.Path(__file__).parent / 'assets' / 'rgb_pytorch.jpg'))
+    img_jpg_nv = decode_image(img_jpg)
+    img_png = read_image(str(pathlib.Path(__file__).parent  / 'assets' / 'rgb_pytorch.png'))
+    assert img_png.ndim == 3 and img_png.numel() > 100
+    img_png_nv = decode_image(img_png)
+
+def smoke_test_torchvision_resnet50_classify() -> None:
     from torchvision.io import read_image
     from torchvision.models import resnet50, ResNet50_Weights
 
-    img = read_image("./test/smoke_test/assets/dog2.jpg")
-
+    img = read_image(str(pathlib.Path(__file__).parent / 'assets' / 'dog2.jpg'))
     # Step 1: Initialize model with the best available weights
     weights = ResNet50_Weights.DEFAULT
     model = resnet50(weights=weights)
@@ -71,12 +68,11 @@ def smoke_test_vision_resnet50_classify() -> None:
     prediction = model(batch).squeeze(0).softmax(0)
     class_id = prediction.argmax().item()
     score = prediction[class_id].item()
-    category_name = weights.meta["categories"][class_id]
-    expected_category = "German shepherd"
-    print(f"{category_name}: {100 * score:.1f}%")
+    category_name = weights.meta['categories'][class_id]
+    expected_category = 'German shepherd'
+    print(f'{category_name}: {100 * score:.1f}%')
     if(category_name != expected_category):
-        print(f"Failed ResNet50 classify  {category_name} Expected: {expected_category}")
-        sys.exit(1)
+        raise RuntimeError(f'Failed ResNet50 classify {category_name} Expected: {expected_category}')
 
 
 def smoke_test_torchaudio() -> None:
@@ -92,14 +88,15 @@ def smoke_test_torchaudio() -> None:
 
 def main() -> None:
     #todo add torch, torchvision and torchaudio tests
-    print(f"torch: {torch.__version__}")
-    print(f"torchvision: {torchvision.__version__}")
-    print(f"torchaudio: {torchaudio.__version__}")
+    print(f'torch: {torch.__version__}')
+    print(f'torchvision: {torchvision.__version__}')
+    print(f'torchaudio: {torchaudio.__version__}')
     smoke_test_cuda()
     smoke_test_conv2d()
-    smoke_test_torchvision()
     smoke_test_torchaudio()
-    smoke_test_vision_resnet50_classify()
+    smoke_test_torchvision()
+    smoke_test_torchvision_read_decode()
+    smoke_test_torchvision_resnet50_classify()
 
-if __name__ == "__main__":
+if __name__ == '__main__':
     main()