Hi all,
I customize a Loss module for seq2seq model. Since inputs are padded to max_seq_len for batch training, so I need a loss that can pass in mask in order to ignore PAD loss.

import torch.nn as nn
import torch
from torch.autograd import Variable

def _assert_no_grad(variable):
	assert not variable.requires_grad, \
		"nn criterions don't compute the gradient w.r.t. targets - please " \
		"mark these variables as volatile or not requiring gradients"

class SequenceLoss(nn.Module):
	def __init__(self, average_cross_timesteps=True, average_cross_batch=True):
		super(SequenceLoss, self).__init__()
		self.average_cross_timesteps = average_cross_timesteps
		self.average_cross_batch = average_cross_batch
		self.loss_func = nn.CrossEntropyLoss(size_average=False)
	
	def forward(self, input, target, weight):
		batch_size = input.size()[0] 
		max_seq_len = input.size()[1]
		_assert_no_grad(target)
		_assert_no_grad(weight)
		num_classes = input.size()[-1]
		logits_flat = input.view(-1, num_classes)
		targets_flat = target.view(-1)
		logits_flat_list = torch.split(logits_flat, 1, dim=0)
		targets_flat_list = torch.split(targets_flat, 1, dim=0)
		tmp = [self.loss_func(logits_flat_list[i], targets_flat_list[i]).data for i in range(batch_size*max_seq_len)]
		
		crossent = Variable(torch.cat(tmp, dim=0), requires_grad=True) #Variable(torch.Tensor(tmp))
		crossent = crossent * weight.view(-1)
		if self.average_cross_timesteps and self.average_cross_batch:
			# char level ppl
			crossent = torch.sum(crossent)
			total_size = torch.sum(weight)
			total_size += 1e-12  # to avoid division by 0 for all-0 weights
			crossent = crossent / total_size
		else:
			sequence_length = input.size()[1]
			crossent = crossent.view(batch_size, sequence_length)
		if self.average_cross_timesteps and not self.average_cross_batch:
			crossent = torch.sum(crossent, 1)
			total_size = torch.sum(weight, 1)
			total_size += 1e-12
			crossent = crossent / total_size
		if not self.average_cross_timesteps and self.average_cross_batch:
			# return crossent = [time_steps]
			crossent = torch.sum(crossent, 0)
			total_size = torch.sum(weight, 0)
			total_size += 1e-12
			crossent = crossent / total_size
		return crossent



And loss doesn’t go down in training process. I don’t know where the problem is. Could anybody give me some advice?