I'm wondering how to implement this with pytorch built-ins. I've got a 3 dimensional input of uints called policy. Most of the entries are zero, and if I were to L1 normalize this I would have a (target) probability distribution.

I've also got the output of a linear layer, called 'logit', with the same shape as 'policy'. I must turn this into a probability distribution by taking the softmax, but only over the entries where policy is non-zero.

The loss is then -sum(log(logit_masked_softmax) * policy_normalized))

I've implemented this manually with the nn.functional module using boolean indexing. The problem is that I want to do this in batches, where a 4 dimensional tensor represents the batches of 3 dimensional inputs. I am convinced that there must be a built-in way to achieve this and it probably is also faster and more numerically stable.

  • $\begingroup$ The same code should work even when batched. Have you tried? Did you get an error? $\endgroup$
    – Jindřich
    Commented Jan 17, 2020 at 10:45
  • $\begingroup$ @Jindřich I tried to implement it but I was having issues with batched softmax to work correctly (no errors, just cant get the dimensions to work). This was after I used a different approach for masking where I replace the zeros with a very large negative number, so that they map to zero after exp. I will update this when my GPU is free $\endgroup$
    – basket
    Commented Jan 17, 2020 at 14:07
  • $\begingroup$ Yes it seems that using the masking approach I orignally used idx = torch.where(p > 0, torch.ones(2, 4, 8, 4).byte(), torch.zeros(2, 4, 8, 4).byte()).type(torch.BoolTensor) and then logit_masked = logit[idx], just returns a list basically, so it doesnt seem possible to compute the correct softmax this way $\endgroup$
    – basket
    Commented Jan 17, 2020 at 14:43

1 Answer 1


I solved this by using torch.where to give the irrelevent entries a really big negative value, so that they vanish after an exp. I also took advantage of the log of CEL trick.

loss = torch.zeros(1).type(dtype)

states = torch.cat([states_[_] for _ in idc[:self.batch_size]]).type(dtype)
scores = torch.tensor([scores_[_] for _ in idc[:self.batch_size]]).type(dtype).view(-1, 1)
policies = torch.cat([policies_[_] for _ in idc[:self.batch_size]]).type(dtype)

values, logits = self.network(states)
value_loss = mse(values, scores)

p = torch.nn.functional.normalize(policies, dim=[1,2,3], p=1)
logit_exp = torch.exp(torch.where(p > 0, logits, inf))
s = torch.log(torch.sum(logit_exp, dim=[1,2,3])).view(-1, 1, 1, 1)
log_q = logits - torch.mul(ones, s)
policy_loss = torch.sum(torch.mul(p, -log_q))/self.batch_size

loss = value_loss + policy_loss
idc = idc[self.batch_size:]

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.