Updated CTCCriterion with variable lengths and batchFirst

1 files changed, 34 insertions, 19 deletions

diff --git a/test/test-all.lua b/test/test-all.lua
index 7486344..7d1d13e 100644
--- a/test/test-all.lua
+++ b/test/test-all.lua
@@ -525,25 +525,40 @@ function nnxtest.TreeNLLCriterion()
 end
 
 function nnxtest.CTCCriterion()
-    local criterion = nn.CTCCriterion()
-    local acts = torch.Tensor({{{0,0,0,0,0}}})
-    local targets = {{1}}
-    mytester:eq(criterion:updateOutput(acts,targets), 1.6094379425049, 0, "CTCCriterion.smallTest")
-    local acts =
-    torch.Tensor({{{1,2,3,4,5}, {6,7,8,9,10}, {11,12,13,14,15}}})
-    local targets = {{3,3}}
-    mytester:eq(criterion:updateOutput(acts,targets), 7.355742931366, 0, "CTCCriterion.mediumTest")
-    local acts = torch.Tensor({{{-5,-4,-3,-2,-1}, {-10,-9,-8,-7,-6}, {-15,-14,-13,-12,-11}}})
-    local targets = {{2,3}}
-    mytester:eq(criterion:updateOutput(acts,targets), 4.938850402832, 0, "CTCCriterion.mediumNegativeTest")
-    local acts =
-    torch.Tensor({
-        {{0,0,0,0,0},{0,0,0,0,0},{0,0,0,0,0}},
-        {{1,2,3,4,5},{6,7,8,9,10},{11,12,13,14,15}},
-        {{-5,-4,-3,-2,-1},{-10,-9,-8,-7,-6},{-15,-14,-13,-12,-11}}
-    })
-    local targets = {{1},{3,3},{2,3}}
-    mytester:eq(criterion:updateOutput(acts,targets), 15.331147670746, 0, "CTCCriterion.batchTest")
+   local criterion = nn.CTCCriterionTest()
+   local acts = torch.Tensor({{{0,0,0,0,0}}}):transpose(1, 2):contiguous() -- input is seqLength x batch x inputDim
+   local targets = {{1}}
+   local sizes = torch.Tensor({1})
+   mytester:eq(criterion:updateOutput(acts, targets, sizes), 1.6094379425049, precision, "CTCCriterion.smallTest")
+   local acts =
+   torch.Tensor({{{1,2,3,4,5}, {6,7,8,9,10}, {11,12,13,14,15}}})
+   local targets = {{3,3}}
+   local sizes = torch.Tensor({3})
+   mytester:eq(criterion:updateOutput(acts, targets, sizes), 7.355742931366, precision, "CTCCriterion.mediumTest")
+   local acts = torch.Tensor({{{-5,-4,-3,-2,-1}, {-10,-9,-8,-7,-6}, {-15,-14,-13,-12,-11}}}):transpose(1, 2):contiguous()
+   local targets = {{2,3}}
+   local sizes = torch.Tensor({3})
+   mytester:eq(criterion:updateOutput(acts, targets, sizes), 4.9388499259949, precision, "CTCCriterion.mediumNegativeTest")
+   local acts =
+   torch.Tensor({
+      {{0,0,0,0,0},{0,0,0,0,0},{0,0,0,0,0}},
+      {{1,2,3,4,5},{6,7,8,9,10},{11,12,13,14,15}},
+      {{-5,-4,-3,-2,-1},{-10,-9,-8,-7,-6},{-15,-14,-13,-12,-11}}
+   }):transpose(1, 2):contiguous()
+   local targets = {{1},{3,3},{2,3}}
+   local sizes = torch.Tensor({1,3,3})
+   mytester:eq(criterion:updateOutput(acts, targets, sizes), 13.904030799866, precision, "CTCCriterion.batchTest")
+   local gradOutputNorm = criterion:updateGradInput(acts, targets, sizes)
+   criterion = nn.CTCCriterionTest(true) -- batchFirst true, input is batch x seqLength x inputDim
+   local batchFirstActs =
+   torch.Tensor({
+      {{0,0,0,0,0},{0,0,0,0,0},{0,0,0,0,0}},
+      {{1,2,3,4,5},{6,7,8,9,10},{11,12,13,14,15}},
+      {{-5,-4,-3,-2,-1},{-10,-9,-8,-7,-6},{-15,-14,-13,-12,-11}}
+   })
+   mytester:eq(criterion:updateOutput(batchFirstActs, targets, sizes), 13.904030799866, precision, "CTCCriterion.batchFirstTest")
+   local gradOutputBatchFirst = criterion:updateGradInput(acts, targets, sizes)
+   mytester:assertTensorEq(gradOutputBatchFirst:transpose(1, 2), gradOutputNorm, precision, "CTCCriterion.gradCheckTest")
 end
 
 local function blur(mean, stdv, size)