Add a test for CTCLoss

SeanNaren · Oct 1, 2020 · fb6c4b4 · fb6c4b4
1 parent 33a97b2
commit fb6c4b4
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 0 deletions.
diff --git a/pytorch_binding/tests/test_cpu.py b/pytorch_binding/tests/test_cpu.py
@@ -69,5 +69,19 @@ def test_empty_label():
     print('CPU_cost: %f' % costs.sum())
 
 
+def test_CTCLoss():
+    probs = torch.FloatTensor([[
+        [0.1, 0.6, 0.1, 0.1, 0.1], [0.1, 0.1, 0.6, 0.1, 0.1]
+    ]]).transpose(0, 1).contiguous()
+    labels = torch.IntTensor([1, 2])
+    label_sizes = torch.IntTensor([2])
+    probs_sizes = torch.IntTensor([2])
+    probs.requires_grad_(True)
+
+    ctc_loss = warp_ctc.CTCLoss()
+    cost = ctc_loss(probs, labels, probs_sizes, label_sizes)
+    cost.backward()
+
+
 if __name__ == '__main__':
     pytest.main([__file__])
diff --git a/pytorch_binding/tests/test_gpu.py b/pytorch_binding/tests/test_gpu.py
@@ -111,5 +111,20 @@ def test_empty_label():
     print(grads.view(grads.size(0) * grads.size(1), grads.size(2)))
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU")
+def test_CTCLoss():
+    probs = torch.FloatTensor([[
+        [0.1, 0.6, 0.1, 0.1, 0.1], [0.1, 0.1, 0.6, 0.1, 0.1]
+    ]]).transpose(0, 1).contiguous().cuda()
+    labels = torch.IntTensor([1, 2])
+    label_sizes = torch.IntTensor([2])
+    probs_sizes = torch.IntTensor([2])
+    probs.requires_grad_(True)
+
+    ctc_loss = warp_ctc.CTCLoss()
+    cost = ctc_loss(probs, labels, probs_sizes, label_sizes)
+    cost.backward()
+
+
 if __name__ == '__main__':
     pytest.main([__file__])