diff --git a/CategoricalEntropy.lua b/CategoricalEntropy.lua new file mode 100644 index 0000000..610494c --- /dev/null +++ b/CategoricalEntropy.lua @@ -0,0 +1,63 @@ +------------------------------------------------------------------------ +--[[ CategoricalEntropy ]]-- +-- Maximize the entropy of a categorical distribution (e.g. softmax ). +-- H(X) = E(-log(p(X)) = -sum(p(X)log(p(X)) +-- where X = 1,...,N and N is the number of categories. +-- A batch with an entropy below minEntropy will be maximized. +-- d H(X=x) p(x) +-- -------- = - ---- - log(p(x)) = -1 - log(p(x)) +-- d p p(x) +------------------------------------------------------------------------ +local CE, parent = torch.class("nn.CategoricalEntropy", "nn.Module") + +function CE:__init(scale, minEntropy) + parent.__init(self) + self.scale = scale or 1 + self.minEntropy = minEntropy + + -- get the P(X) using the batch as a prior + self.module = nn.Sequential() + self.module:add(nn.Sum(1)) -- sum categorical probabilities over batch + self._mul = nn.MulConstant(1) + self.module:add(self._mul) -- make them sum to one (i.e. probabilities) + + -- get entropy H(X) + local concat = nn.ConcatTable() + concat:add(nn.Identity()) -- p(X) + local seq = nn.Sequential() + seq:add(nn.AddConstant(0.000001)) -- prevent log(0) = nan errors + seq:add(nn.Log()) + concat:add(seq) + self.module:add(concat) -- log(p(x)) + self.module:add(nn.CMulTable()) -- p(x)log(p(x)) + self.module:add(nn.Sum()) -- sum(p(x)log(p(x))) + self.module:add(nn.MulConstant(-1)) -- H(x) + + self.modules = {self.module} + + self.minusOne = torch.Tensor{-self.scale} -- gradient descent on maximization + self.sizeAverage = true +end + +function CE:updateOutput(input) + assert(input:dim() == 2, "CategoricalEntropy only works with batches") + self.output:set(input) + return self.output +end + +function CE:updateGradInput(input, gradOutput, scale) + assert(input:dim() == 2, "CategoricalEntropy only works with batches") + self.gradInput:resizeAs(input):copy(gradOutput) + + self._mul.constant_scalar = 1/input:sum() -- sum to one + self.entropy = self.module:updateOutput(input)[1] + if (not self.minEntropy) or (self.entropy < self.minEntropy) then + local gradEntropy = self.module:updateGradInput(input, self.minusOne, scale) + if self.sizeAverage then + gradEntropy:div(input:size(1)) + end + self.gradInput:add(gradEntropy) + end + + return self.gradInput +end diff --git a/init.lua b/init.lua index 4d3d174..c596aee 100644 --- a/init.lua +++ b/init.lua @@ -70,6 +70,7 @@ torch.include('dpnn', 'Clip.lua') torch.include('dpnn', 'SpatialUniformCrop.lua') torch.include('dpnn', 'SpatialGlimpse.lua') torch.include('dpnn', 'ArgMax.lua') +torch.include('dpnn', 'CategoricalEntropy.lua') -- REINFORCE torch.include('dpnn', 'Reinforce.lua') diff --git a/test/test.lua b/test/test.lua index 9c65f36..0b1fcf8 100644 --- a/test/test.lua +++ b/test/test.lua @@ -705,6 +705,28 @@ function dpnntest.ArgMax() mytester:assertTensorEq(gradInput, input:clone():zero(), 0.000001, "ArgMax gradInput not asLong err") end +function dpnntest.CategoricalEntropy() + local inputSize = 5 + local batchSize = 10 + local minEntropy = 12 + local input_ = torch.randn(batchSize, inputSize) + local input = nn.SoftMax():updateOutput(input_) + local gradOutput = torch.Tensor(batchSize, inputSize):zero() + local ce = nn.CategoricalEntropy() + local output = ce:forward(input) + mytester:assertTensorEq(input, output, 0.0000001, "CategoricalEntropy output err") + local gradInput = ce:backward(input, gradOutput) + local output2 = input:sum(1)[1] + output2:div(output2:sum()) + local log2 = torch.log(output2 + 0.000001) + local entropy2 = -output2:cmul(log2):sum() + mytester:assert(math.abs(ce.entropy - entropy2) < 0.000001, "CategoricalEntropy entropy err") + local gradEntropy2 = log2:add(1) -- -1*(-1 - log(p(x))) = 1 + log(p(x)) + gradEntropy2:div(input:sum()) + local gradInput2 = gradEntropy2:div(batchSize):view(1,inputSize):expandAs(input) + mytester:assertTensorEq(gradInput2, gradInput, 0.000001, "CategoricalEntropy gradInput err") +end + function dpnnbigtest.Reinforce() -- let us try to reinforce an mlp to learn a simple distribution local n = 10