Basic CNN RNN Test (pytorch#564)

* rnn cnn smoke test

* only for cuda11 first

* rm cuda11 condition

* some refactor

* add missing file

* fix wording error

* minor change

* share the smoke tests for linux

* Update check_binary.sh

Co-authored-by: Nikita Shulga <[email protected]>

Author: mszhanyi

check_binary.sh

@@ -362,6 +362,13 @@ if [[ "$DESIRED_CUDA" != 'cpu' && "$DESIRED_CUDA" != *"rocm"* ]]; then
 
     echo "Checking that CuDNN is available"
     python -c 'import torch; exit(0 if torch.backends.cudnn.is_available() else 1)'
+
+    echo "Checking that basic RNN works"
+    python ${TEST_CODE_DIR}/rnn_smoke.py
+
+    echo "Checking that basic CNN works"
+    python ${TEST_CODE_DIR}/cnn_smoke.py
+
     popd
   fi # if libtorch
 fi # if cuda

test_example_code/cnn_smoke.py

@@ -0,0 +1,36 @@
+r"""
+It's used to check basic rnn features with cuda.
+For example, it would throw exception if some components are missing
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+
+class SimpleCNN(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv = nn.Conv2d(1, 1, 3)
+        self.pool = nn.MaxPool2d(2, 2)
+
+    def forward(self, inputs):
+        output = self.pool(F.relu(self.conv(inputs)))
+        output = output.view(1)
+        return output
+
+# Mock one infer
+device = torch.device("cuda:0")
+net = SimpleCNN().to(device)
+net_inputs = torch.rand((1, 1, 5, 5), device=device)
+outputs = net(net_inputs)
+print(outputs)
+
+criterion = nn.MSELoss()
+optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.1)
+
+# Mock one step training
+label = torch.full((1,), 1.0, dtype=torch.float, device=device)
+loss = criterion(outputs, label)
+loss.backward()
+optimizer.step()

test_example_code/rnn_smoke.py

@@ -0,0 +1,13 @@
+r"""
+It's used to check basic rnn features with cuda.
+For example, it would throw exception if missing some components are missing
+"""
+
+import torch
+import torch.nn as nn
+
+device = torch.device("cuda:0")
+rnn = nn.RNN(10, 20, 2).to(device)
+inputs = torch.randn(5, 3, 10).to(device)
+h0 = torch.randn(2, 3, 20).to(device)
+output, hn = rnn(inputs, h0)

windows/internal/smoke_test.bat

@@ -146,6 +146,14 @@ echo Checking that CuDNN is available
 python -c "import torch; exit(0 if torch.backends.cudnn.is_available() else 1)"
 if ERRORLEVEL 1 exit /b 1
 
+echo Checking that basic RNN works
+python %BUILDER_ROOT%\test_example_code\rnn_smoke.py
+if ERRORLEVEL 1 exit /b 1
+
+echo Checking that basic CNN works
+python %BUILDER_ROOT%\test_example_code\cnn_smoke.py
+if ERRORLEVEL 1 exit /b 1
+
 goto end
 
 :libtorch