Убыток не уменьшается для Ptorch LSTM

Question

Я новичок в pytorch и ищу вашу помощь в реализации lstm. У меня есть однослойный LSTM, за которым следуют полностью подключенный слой и сигмоид (реализующий трассировку Deep Knowledge).

Примеры входных и выходных пар следующие: input =

[[1, 0,0, 0],

[0, 1, 0, 0],

[0, 0, 1, 0],

[0, 0, 0,1]]

output = [[1, 0],

[0, 0],

[0, 1],

[-1, -1]]

У меня есть следующий код для LSTM и я рассчитываю вычислить двоичную перекрестную энтропию как потерю. Однако потери не уменьшаются в течение эпох. Я проследил несколько блогов, чтобы реализовать это, и я думаю, что это правильно. Что-то не так с кодом, который у меня есть? Было бы здорово, если бы вы могли потратить пару минут на просмотр кода и помочь подсказать, если с ним что-то не так.

    class BuildModel(nn.Module):
        on_gpu = False
        def __init__(self, num_skills, batch_size = 2, lstm_units = 200):
            super(BuildModel, self).__init__()
            self.lstm_units = lstm_units
            self.batch_size = batch_size
            self.output_dim = num_skills
            self.input_dim = num_skills * 2
            self.seq_len = seq_len
            self.__build_model()

        def __build_model(self):
            self.lstm = nn.LSTM(
                    input_size = self.input_dim,
                    hidden_size = self.lstm_units,
                    num_layers = 1,
                    batch_first = True,
                    )
            self.hidden_to_skills = nn.Linear(self.lstm_units, self.output_dim)

        def init_hidden(self):
            hidden_a = torch.randn(1, self.batch_size, self.lstm_units)
            hidden_b = torch.randn(1, self.batch_size, self.lstm_units)

            if self.on_gpu:
                hidden_a = hidden_a.cuda()
                hidden_b = hidden_b.cuda()

            hidden_a = Variable(hidden_a, requires_grad=True)
            hidden_b = Variable(hidden_b, requires_grad=True)
            return (hidden_a, hidden_b)

        def forward(self, X, X_lengths):
            self.hidden = self.init_hidden()

            batch_size, seq_len, _ = X.size()

            ## X is of shape batch_size, sequence_length, num_skills*2
            X = torch.nn.utils.rnn.pack_padded_sequence(X, X_lengths, batch_first=True, enforce_sorted=False)

            X, self.hidden = self.lstm(X, self.hidden)

            X, _ = torch.nn.utils.rnn.pad_packed_sequence(X, batch_first=True)


            # Transfer data from (batch_size, seq+len, lstm_units) --> (batch_size * seq_len, lstm_units)
            X = X.contiguous()
            X = X.view(-1, X.shape[2])

            X = self.hidden_to_skills(X)

            X = torch.nn.functional.sigmoid(X)
            # return the predictions
            return X

        def loss(self, Y_hat, Y, threshold, seq_len):
            #flatten labels
            Y = Y.view(-1)

            Y_hat = Y_hat.view(-1, seq_len * self.output_dim * self.batch_size)

            mask = (Y > -1).float()
            mask_long = (Y > -1).long()
            nb_tokens = int(torch.sum(mask).item())

            Y_hat = Y_hat[range(Y_hat.shape[0])] * mask
            Y_hat = Y_hat.reshape(-1)
            Y = Y[range(Y.shape[0])] * mask_long
            Y = Y.float()
            Y_hat = (Y_hat > 0.0).float() * 1.0
            loss = torch.nn.BCELoss()
            ce_loss = loss(Y_hat, Y)

            return Variable(ce_loss, requires_grad = True)

### For training

learning_rate = .1
total_step = len(loader)
model = BuildModel(num_skills, 2, 200)
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)


for epoch in range(num_epochs):
      overall_loss = 0.0
      for i, (X, y, lengths, seq_len) in enumerate(loader):
          X = torch.from_numpy(X).float()
          y = torch.from_numpy(y).long()

          outputs = model(X, lengths)
          loss = model.loss(outputs, y, 0.5, seq_len)
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()