MaskZeroCriterion v2

Author: Nicholas Leonard

AbstractSequencerCriterion.lua

@@ -0,0 +1,49 @@
+------------------------------------------------------------------------
+--[[ AbstractSequencerCriterion ]]--
+-- Inherited by SequencerCriterion and RepeaterCriterion
+-- WARNING : assumes that the decorated criterion is stateless, i.e.
+-- the backward doesn't need to be preceded by a commensurate forward.
+------------------------------------------------------------------------
+local AbstractSequencerCriterion, parent = torch.class('nn.AbstractSequencerCriterion', 'nn.Criterion')
+
+function AbstractSequencerCriterion:__init(criterion, sizeAverage)
+   parent.__init(self)
+   self.criterion = criterion
+   if torch.isTypeOf(criterion, 'nn.ModuleCriterion') then
+      error(torch.type(self).." shouldn't decorate a ModuleCriterion. "..
+         "Instead, try the other way around : "..
+         "ModuleCriterion decorates a ".. torch.type(self) .. ". "..
+         "Its modules can also be similarly decorated with a Sequencer.")
+   end
+   if sizeAverage ~= nil then
+      self.sizeAverage = sizeAverage
+   else
+      self.sizeAverage = false
+   end
+   self.clones = {}
+end
+
+function AbstractSequencerCriterion:getStepCriterion(step)
+   assert(step, "expecting step at arg 1")
+   local criterion = self.clones[step]
+   if not criterion then
+      criterion = self.criterion:clone()
+      self.clones[step] = criterion
+   end
+   return criterion
+end
+
+function AbstractSequencerCriterion:setZeroMask(zeroMask)
+   if zeroMask == false then
+      for k,stepcriterion in pairs(self.clones) do
+         stepcriterion:setZeroMask(zeroMask)
+      end
+   else
+      assert(zeroMask:dim() >= 2, "Expecting dim >= 2 for zeroMask. For example, seqlen x batchsize")
+      for step=1,zeroMask:size(1) do
+         local stepcriterion = self:getStepCriterion(step)
+         stepcriterion:setZeroMask(zeroMask[step])
+      end
+   end
+end
+

CMakeLists.txt

@@ -15,6 +15,7 @@ SET(luasrc
   init.lua
   AbstractRecurrent.lua
   AbstractSequencer.lua
+  AbstractSequencerCriterion.lua
   BiSequencer.lua
   BiSequencerLM.lua
   CopyGrad.lua
@@ -33,7 +34,6 @@ SET(luasrc
   Padding.lua
   Recurrence.lua
   RecurrentAttention.lua
-  recursiveUtils.lua
   Recursor.lua
   Repeater.lua
   RepeaterCriterion.lua

Criterion.lua

@@ -2,3 +2,15 @@ local Criterion = nn.Criterion
 
 Criterion.toBatch = nn.Module.toBatch
 Criterion.fromBatch = nn.Module.fromBatch
+
+
+function Criterion:setZeroMask(zeroMask)
+   if self.criterions then
+      for i, criterion in ipairs(self.criterions) do
+         criterion:setZeroMask(zeroMask)
+      end
+   end
+   if self.criterion then
+   	  self.criterion:setZeroMask(zeroMask)
+   end
+end

MaskZeroCriterion.lua

@@ -5,128 +5,84 @@
 ------------------------------------------------------------------------
 local MaskZeroCriterion, parent = torch.class("nn.MaskZeroCriterion", "nn.Criterion")
 
-function MaskZeroCriterion:__init(criterion, nInputDim)
+function MaskZeroCriterion:__init(criterion)
    parent.__init(self)
    self.criterion = criterion
    assert(torch.isTypeOf(criterion, 'nn.Criterion'))
-   assert(torch.type(nInputDim) == 'number', 'Expecting nInputDim number at arg 2')
-   self.nInputDim = nInputDim
+   self.v2 = true
 end
 
-function MaskZeroCriterion:recursiveGetFirst(input)
-   if torch.type(input) == 'table' then
-      return self:recursiveGetFirst(input[1])
-   else
-      assert(torch.isTensor(input))
-      return input
-   end
-end
-
-function MaskZeroCriterion:recursiveMask(dst, src, mask)
-   if torch.type(src) == 'table' then
-      dst = torch.type(dst) == 'table' and dst or {}
-      for k,v in ipairs(src) do
-         dst[k] = self:recursiveMask(dst[k], v, mask)
+function MaskZeroCriterion:updateOutput(input, target)
+   if self.v2 then
+      assert(self.zeroMask ~= nil, "MaskZeroCriterion expecting zeroMask tensor or false")
+      if self.zeroMask == false then
+         self.output = self.criterion:updateOutput(input, target)
+         return self.output
       end
-   else
-      assert(torch.isTensor(src))
-      dst = torch.isTensor(dst) and dst or src.new()
-   	
-      dst:index(src, 1, mask)
+      assert(self.zeroMask:dim() == 1, "MaskZeroCriterion expecting zeroMask of size batchsize")
+   else -- backwards compat
+      self.zeroMask = nn.utils.getZeroMaskBatch(input, self.zeroMask)
    end
-   return dst
-end
 
-function MaskZeroCriterion:updateOutput(input, target)   
-   -- recurrent module input is always the first one
-   local rmi = self:recursiveGetFirst(input):contiguous()
-   if rmi:dim() == self.nInputDim then
-      error("does not support online (i.e. non-batch) mode")
-   elseif rmi:dim() - 1 == self.nInputDim then
-      rmi = rmi:view(rmi:size(1), -1) -- collapse non-batch dims
+   self.isEmptyBatch = (self.zeroMask:sum() == self.zeroMask:nElement())
+   if self.isEmptyBatch then
+      self.output = 0
    else
-      error("nInputDim error: "..rmi:dim()..", "..self.nInputDim)
-   end
-   
-   -- build mask
-   local vectorDim = rmi:dim() 
-   self._zeroMask = self._zeroMask or rmi.new()
-   self._zeroMask:norm(rmi, 2, vectorDim)
-   local zeroMask = self._zeroMask
-   if torch.isTypeOf(zeroMask, 'torch.CudaTensor') or
-         torch.isTypeOf(zeroMask, 'torch.ClTensor') then
-      self.__zeroMask = self.__zeroMask or torch.FloatTensor()
-      self.__zeroMask:resize(self._zeroMask:size()):copy(self._zeroMask)
-      zeroMask = self._zeroMask
-   end
-  
-   self.zeroMask = self.zeroMask or torch.LongTensor()
-   self.zeroMask:resize(self._zeroMask:size(1)):zero()
-   
-   local i, j = 0, 0
-   zeroMask:apply(function(norm)
-      i = i + 1
-      if norm ~= 0 then
-         j = j + 1
-         self.zeroMask[j] = i
-      end
-   end)
-   self.zeroMask:resize(j)
-   
-   if j > 0 then
-      self.input = self:recursiveMask(self.input, input, self.zeroMask)
-      self.target = self:recursiveMask(self.target, target, self.zeroMask)
-      
+      -- e.g. 0,1,0 -> 1,0,1
+      self._oneMask = self._oneMask or self.zeroMask.new()
+      self._oneMask:lt(self.zeroMask, 1)
+      -- 1,0,1 -> 1,3
+      self._indices = self._indices or torch.LongTensor()
+      self._range = self._range or torch.LongTensor()
+      self._range:range(1,self._oneMask:nElement())
+      self._indices:maskedSelect(self._range, self._oneMask)
+      -- indexSelect the input
+      self.input = nn.utils.recursiveIndex(self.input, input, 1, self._indices)
+      self.target = nn.utils.recursiveIndex(self.target, target, 1, self._indices)
+
       -- forward through decorated criterion
       self.output = self.criterion:updateOutput(self.input, self.target)
-   else
-      -- when all samples are masked, then loss is zero (issue 128)
-      self.output = 0
    end
-   
+
    return self.output
 end
 
-function MaskZeroCriterion:recursiveMaskGradInput(dst, mask, src, input)
-   if torch.type(input) == 'table' then
-      dst = (torch.type(dst) == 'table') and dst or {dst}
-      src = (torch.type(src) == 'table') and src or {src}
-      for key,_ in pairs(input) do
-         dst[key] = self:recursiveMaskGradInput(dst[key], mask, src[key], input[key])
-      end
-      for i=#input+1,#dst do
-         dst[i] = nil
-      end
-   elseif torch.isTensor(input) then
-      dst = torch.isTensor(dst) and dst or input.new()
-      dst:resizeAs(input):zero()
-      if mask:nElement() > 0 then
-         assert(src)
-         dst:indexCopy(1, mask, src)
-      end
-   else
-      error("expecting nested tensors or tables. Got "..
-            torch.type(dst).." and "..torch.type(input).." instead")
+function MaskZeroCriterion:updateGradInput(input, target)
+   if self.zeroMask == false then
+      self.gradInput = self.criterion:updateGradInput(input, target)
+      return self.gradInput
    end
-   return dst
-end
 
-function MaskZeroCriterion:updateGradInput(input, target)
-   if self.zeroMask:nElement() > 0 then
+   self._gradInput = nn.utils.recursiveResizeAs(self._gradInput, input)
+   nn.utils.recursiveFill(self._gradInput, 0)
+
+   if not self.isEmptyBatch then
       assert(self.input and self.target)
-      self._gradInput = self.criterion:updateGradInput(self.input, self.target)
+      local gradInput = self.criterion:updateGradInput(self.input, self.target)
+      nn.utils.recursiveIndexCopy(self._gradInput, 1, self._indices, gradInput)
    end
-   self.gradInput = self:recursiveMaskGradInput(self.gradInput, self.zeroMask, self._gradInput, input)
+
+   self.gradInput = self._gradInput
    return self.gradInput
 end
 
-function MaskZeroCriterion:type(type, ...)
+function MaskZeroCriterion:clearState()
    self.zeroMask = nil
-   self._zeroMask = nil
-   self.__zeroMask = nil
+   self._oneMask = nil
+   self._range = nil
+   self._indices = nil
    self.input = nil
    self.target = nil
+   self.output = nil
+   self.gradInput = nil
    self._gradInput = nil
-   
+end
+
+function MaskZeroCriterion:type(type, ...)
+   self:clearState()
    return parent.type(self, type, ...)
 end
+
+function MaskZeroCriterion:setZeroMask(zeroMask)
+   self.zeroMask = zeroMask
+end

README.md

@@ -1192,20 +1192,25 @@ This lookup table makes it possible to pad sequences with different lengths in t
 
 <a name='rnn.MaskZeroCriterion'></a>
 ## MaskZeroCriterion ##
-This criterion zeroes the `err` and `gradInput` rows of the decorated criterion
-for commensurate `input` rows which are tensors of zeros.
+
+This criterion ignores samples (rows in the `input` and `target` tensors)
+where the `zeroMask` ByteTensor passed to `MaskZeroCriterion:setZeroMask(zeroMask)` is 1.
+This criterion only supports batch-mode.
 
 ```lua
-mzc = nn.MaskZeroCriterion(criterion, nInputDim)
+batchsize = 3
+zeroMask = torch.ByteTensor(batchsize):zero()
+zeroMask[2] = 1 -- the 2nd sample in batch is ignored
+mzc = nn.MaskZeroCriterion(criterion)
+mzc:setZeroMask(zeroMask)
+loss = mzc:forward(input, target)
+gradInput = mzc:backward(input, target)
+assert(gradInput[2]:sum() == 0)
 ```
 
-The `gradInput` Tensor (or table thereof) of the decorated `criterion`
-will have each row (samples) zeroed when the commensurate row of the `input`
-is a tensor of zeros. The `err` will also disregard such zero rows.
-
-The `nInputDim` argument must specify the number of non-batch dims
-in the first Tensor of the `input`. In the case of an `input` table,
-the first Tensor is the first one encountered when doing a depth-first search.
+In the above example, the second row of the `gradInput` Tensor is zero.
+This is because the commensurate row in the `zeroMask` is a one.
+The call to `forward` also disregards the second sample in measuring the `loss`.
 
 This decorator makes it possible to pad sequences with different lengths in the same batch with zero vectors.
 

RepeaterCriterion.lua

@@ -4,59 +4,57 @@
 -- same target (the target is repeated).
 -- Useful for nn.Repeater and nn.Sequencer.
 ------------------------------------------------------------------------
-assert(not nn.RepeaterCriterion, "update nnx package : luarocks install nnx")
-local RepeaterCriterion, parent = torch.class('nn.RepeaterCriterion', 'nn.Criterion')
+local RepeaterCriterion, parent = torch.class('nn.RepeaterCriterion', 'nn.AbstractSequencerCriterion')
 
-function RepeaterCriterion:__init(criterion)
-   parent.__init(self)
-   self.criterion = criterion
-   self.gradInput = {}
-   self.clones = {}
-end
-
-RepeaterCriterion.getStepCriterion = nn.SequencerCriterion.getStepCriterion
-
-function RepeaterCriterion:forward(input, target)
+function RepeaterCriterion:updateOutput(input, target)
    self.output = 0
-   local nStep
+   local seqlen
    if torch.isTensor(input) then
-      nStep = input:size(1)
+      seqlen = input:size(1)
    else
-      nStep = #input
+      seqlen = #input
    end
 
-   
-   for i=1,nStep do
+   for i=1,seqlen do
       local criterion = self:getStepCriterion(i)
       self.output = self.output + criterion:forward(input[i], target)
    end
-   
+
+
+   if self.sizeAverage then
+      self.output = self.output / seqlen
+   end
+
    return self.output
 end
 
-function RepeaterCriterion:backward(input, target)
+function RepeaterCriterion:updateGradInput(input, target)
    self.gradInput = {}
    if torch.isTensor(input) then
-      nStep = input:size(1)
+      seqlen = input:size(1)
    else
-      nStep = #input
+      seqlen = #input
    end
-   
+
    local tableGradInput = {}
-   for i=1,nStep do
+   for i=1,seqlen do
       local criterion = self:getStepCriterion(i)
       tableGradInput[i] = criterion:backward(input[i], target)
    end
-   
+
+   if self.sizeAverage then
+      nn.utils.recursiveDiv(tableGradInput[i], seqlen)
+   end
+
    if torch.isTensor(input) then
       self.gradInput = tableGradInput[1].new()
-      self.gradInput:resize(nStep, unpack(tableGradInput[1]:size():totable()))
-      for step=1,nStep do
+      self.gradInput:resize(seqlen, unpack(tableGradInput[1]:size():totable()))
+      for step=1,seqlen do
          self.gradInput[step]:copy(tableGradInput[step])
       end
    else
       self.gradInput = tableGradInput
    end
-   
+
    return self.gradInput
 end