Smoke test modifications

Author: atalman

test/smoke_test/smoke_test.py

@@ -4,11 +4,11 @@
 import torchvision
 import torchaudio
 
-def smoke_test_cuda() -> None:
-    gpu_arch_ver = os.getenv('GPU_ARCH_VER')
-    gpu_arch_type = os.getenv('GPU_ARCH_TYPE')
-    is_cuda_system = gpu_arch_type == "cuda"
+gpu_arch_ver = os.getenv('GPU_ARCH_VER')
+gpu_arch_type = os.getenv('GPU_ARCH_TYPE')
+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)
@@ -21,6 +21,25 @@ def smoke_test_cuda() -> None:
         #todo add cudnn version validation
         print(f"torch cudnn: {torch.backends.cudnn.version()}")
 
+def smoke_test_conv2d() -> None:
+    import torch.nn as nn
+    print(f"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
+    m = nn.Conv2d(16, 33, (3, 5), stride=(2, 1), padding=(4, 2))
+    # non-square kernels and unequal stride and with padding and dilation
+    m = nn.Conv2d(16, 33, (3, 5), stride=(2, 1), padding=(4, 2), dilation=(3, 1))
+    input = torch.randn(20, 16, 50, 100)
+    output = m(input)
+    if(is_cuda_system):
+        print(f"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
@@ -53,7 +72,7 @@ def smoke_test_vision_resnet50_classify() -> None:
     class_id = prediction.argmax().item()
     score = prediction[class_id].item()
     category_name = weights.meta["categories"][class_id]
-    expected_category = "German sheppard"
+    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}")
@@ -77,6 +96,7 @@ def main() -> None:
     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()