# batched CrossEntropyLoss in pytorch

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.

• The same code should work even when batched. Have you tried? Did you get an error? Jan 17 '20 at 10:45
• @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 Jan 17 '20 at 14:07
• 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 Jan 17 '20 at 14:43

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
loss.backward()
optimizer.step()