Понимание обратного механизма LSTMCell в Pytorch

Question

Я хочу подключиться к обратному проходу функции LSTMCell в pytorch, поэтому на этапе инициализации я делаю следующее (num_layers = 4, hidden_size = 1000, input_size = 1000):

self.layers = nn.ModuleList([
        LSTMCell(
            input_size=input_size,
            hidden_size=hidden_size,
        )
        for layer in range(num_layers)
    ])

for l in self.layers:
    l.register_backward_hook(backward_hook)

В прямом проходе я просто повторяю LSTMCell по длине последовательности и num_layers следующим образом:

for j in range(seqlen):            
    input = #some tensor of size (batch_size, input_size)
    for i, rnn in enumerate(self.layers):
        # recurrent cell
        hidden, cell = rnn(input, (prev_hiddens[i], prev_cells[i]))

Если вход имеет размер (batch_size, input_size), prev_hiddens[i] имеет размер (batch_size, hidden_size), prev_cells[i] имеет размер (batch_size, hidden_size).

В backward_hook я печатаю размер тензоров, которые вводятся в эту функцию:

def backward_hook(module, grad_input, grad_output):
    for grad in grad_output:
        print ("grad_output {}".format(grad))

    for grad in grad_input:
         print ("grad_input.size () {}".format(grad.size()))

В качестве результата впервые backward_hook называется, например:

[A] Для grad_output я получаю 2 тензора, среди которых второй тензор равен None. Это понятно, потому что в обратной фазе у нас есть градиент внутренних состояний (c) и градиент выхода (h). Последняя итерация во временном измерении не имеет скрытого будущего, поэтому ее градиент равен None.

[B] Для grad_input я получаю 5 тензоров (batch_size = 9):

grad_input.size () torch.Size([9, 4000])
grad_input.size () torch.Size([9, 4000])
grad_input.size () torch.Size([9, 1000])
grad_input.size () torch.Size([4000])
grad_input.size () torch.Size([4000])

Мои вопросы:

(1) Правильно ли мое понимание [A]?

(2) Как мне интерпретировать 5 тензоров из кортежа grad_input? Я думал, что должно быть только 3, так как есть только 3 входа в LSTMCell forward ()?

Спасибо

Answer 1

Ваше понимание grad_input и grad_output неверно.Я пытаюсь объяснить это на более простом примере.

def backward_hook(module, grad_input, grad_output):
    for grad in grad_output:
        print ("grad_output.size {}".format(grad.size()))

    for grad in grad_input:
        if grad is None:
            print('None')
        else:
            print ("grad_input.size: {}".format(grad.size()))
    print()

model = nn.Linear(10, 20)
model.register_backward_hook(backward_hook)

input = torch.randn(8, 3, 10)
Y = torch.randn(8, 3, 20)

Y_pred = []
for i in range(input.size(1)):
    out = model(input[:, i])
    Y_pred.append(out)

loss = torch.norm(Y - torch.stack(Y_pred, dim=1), 2)
loss.backward()

Вывод:

grad_output.size torch.Size([8, 20])
grad_input.size: torch.Size([8, 20])
None
grad_input.size: torch.Size([10, 20])

grad_output.size torch.Size([8, 20])
grad_input.size: torch.Size([8, 20])
None
grad_input.size: torch.Size([10, 20])

grad_output.size torch.Size([8, 20])
grad_input.size: torch.Size([8, 20])
None
grad_input.size: torch.Size([10, 20])

Объяснение

• grad_output: градиент потерь по выходу слоя, Y_pred.

• grad_input: градиент потерь по входам слоя.Для слоя Linear входными данными являются тензор input и weight и bias.

Итак, в выводе вы видите:

grad_input.size: torch.Size([8, 20])  # for the `bias`
None                                  # for the `input`
grad_input.size: torch.Size([10, 20]) # for the `weight`

---

Слой Linear в PyTorch использует LinearFunction, который выглядит следующим образом.

class LinearFunction(Function):

    # Note that both forward and backward are @staticmethods
    @staticmethod
    # bias is an optional argument
    def forward(ctx, input, weight, bias=None):
        ctx.save_for_backward(input, weight, bias)
        output = input.mm(weight.t())
        if bias is not None:
            output += bias.unsqueeze(0).expand_as(output)
        return output

    # This function has only a single output, so it gets only one gradient
    @staticmethod
    def backward(ctx, grad_output):
        # This is a pattern that is very convenient - at the top of backward
        # unpack saved_tensors and initialize all gradients w.r.t. inputs to
        # None. Thanks to the fact that additional trailing Nones are
        # ignored, the return statement is simple even when the function has
        # optional inputs.
        input, weight, bias = ctx.saved_tensors
        grad_input = grad_weight = grad_bias = None

        # These needs_input_grad checks are optional and there only to
        # improve efficiency. If you want to make your code simpler, you can
        # skip them. Returning gradients for inputs that don't require it is
        # not an error.
        if ctx.needs_input_grad[0]:
            grad_input = grad_output.mm(weight)
        if ctx.needs_input_grad[1]:
            grad_weight = grad_output.t().mm(input)
        if bias is not None and ctx.needs_input_grad[2]:
            grad_bias = grad_output.sum(0).squeeze(0)

        return grad_input, grad_weight, grad_bias

---

Для LSTM существует четыре набора весовых параметров.

weight_ih_l0
weight_hh_l0
bias_ih_l0
bias_hh_l0

Итак, в вашем случае grad_input будет набором из 5 тензоров.И как вы упомянули, grad_output - это два тензора.