merged

Author: nicholas-leonard

VRClassReward.lua

@@ -22,19 +22,27 @@ function VRClassReward:updateOutput(input, target)
    assert(torch.type(input) == 'table')
    local input = self:toBatch(input[1], 1)
    self._maxVal = self._maxVal or input.new()
-   self._maxIdx = self._maxIdx or torch.type(input) == 'torch.CudaTensor' and input.new() or torch.LongTensor()
+   self._maxIdx = self._maxIdx or torch.type(input) == 'torch.CudaTensor' and torch.CudaLongTensor() or torch.LongTensor()
 
    -- max class value is class prediction
    self._maxVal:max(self._maxIdx, input, 2)
-   if torch.type(self._maxIdx) ~= torch.type(target) then
-      self._target = self._target or self._maxIdx.new()
+   
+   -- reward = scale when correctly classified
+   local maxIdx = self._maxIdx
+   if torch.type(self._maxIdx) == 'torch.CudaLongTensor' then
+      self.__maxIdx = self.__maxIdx or torch.CudaTensor()
+      self.__maxIdx:resize(maxIdx:size()):copy(maxIdx)
+      maxIdx = self.__maxIdx
+   end
+   
+   if torch.type(maxIdx) ~= torch.type(target) then
+      self._target = self._target or maxIdx.new()
       self._target:resize(target:size()):copy(target)
       target = self._target
    end
 
-   -- reward = scale when correctly classified
-   self._reward = self._maxIdx.new()
-   self._reward:eq(self._maxIdx, target)
+   self._reward = self._reward or maxIdx.new()
+   self._reward:eq(maxIdx, target)
    self.reward = self.reward or input.new()
    self.reward:resize(self._reward:size(1)):copy(self._reward)
    self.reward:mul(self.scale)
@@ -66,6 +74,7 @@ function VRClassReward:updateGradInput(inputTable, target)
    self.gradInput[1] = self:fromBatch(self.gradInput[1], 1)
 
    -- learn the baseline reward
+   self.criterion:forward(baseline, self.reward)
    self.gradInput[2] = self.criterion:backward(baseline, self.reward)
    self.gradInput[2] = self:fromBatch(self.gradInput[2], 1)
    return self.gradInput
@@ -74,6 +83,7 @@ end
 function VRClassReward:type(type)
    self._maxVal = nil
    self._maxIdx = nil
+   self.__maxIdx = nil
    self._target = nil
    local module = self.module
    self.module = nil

test/test.lua

@@ -806,19 +806,34 @@ function dpnntest.ReinforceCategorical()
 end
 
 function dpnntest.VRClassReward()
-   local input = {torch.randn(13,10), torch.randn(13,1)}
+   local input = {torch.randn(13,10):float(), torch.randn(13,1):float()}
    local target = torch.IntTensor(13):random(1,10)
-   local rf = nn.Reinforce()
-   local vrc = nn.VRClassReward(rf)
+   local rf = nn.Reinforce():float()
+   local vrc = nn.VRClassReward(rf):float()
    local err = vrc:forward(input, target)
    local gradInput = vrc:backward(input, target)
    local val, idx = input[1]:max(2)
-   local reward = torch.eq(idx:select(2,1):int(), target):double()
+   local reward = torch.eq(idx:select(2,1):int(), target):float()
    local err2 = -reward:mean()
    mytester:assert(err == err2, "VRClassReward forward err")
-   local gradInput2 = nn.MSECriterion():backward(input[2], reward)
+   local gradInput2 = nn.MSECriterion():float():backward(input[2], reward)
    mytester:assertTensorEq(gradInput[2], gradInput2, 0.000001, "VRClassReward backward baseline err")
-   mytester:assertTensorEq(gradInput[1], input[1]:zero(), 0.000001, "VRClassReward backward class err")
+   mytester:assert(math.abs(gradInput[1]:sum()) < 0.000001, "VRClassReward backward class err")
+   
+   if pcall(function() require 'cunn' end) then
+      local gradInput = {gradInput[1], gradInput[2]}
+      input[1], input[2] = input[1]:cuda(), input[2]:cuda()
+      target = target:cuda()
+      rf:cuda()
+      vrc:cuda()
+      
+      local err2 = vrc:forward(input, target)
+      local gradInput2 = vrc:backward(input, target)
+      
+      mytester:assert(math.abs(err - err2) < 0.000001, "VRClassReward forward cuda err")
+      mytester:assertTensorEq(gradInput[2], gradInput2[2]:float(), 0.000001, "VRClassReward backward baseline cuda err")
+      mytester:assertTensorEq(gradInput[1], gradInput2[1]:float(), 0.000001, "VRClassReward backward class cuda err")
+   end
 end
 
 function dpnntest.Clip()