Understanding the backward mechanism of LSTMCell in Pytorch - neural-network

Question

I want to hook into the backward pass of a LSTMCell function in pytorch so in the initialization pass I do the following (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)
In the forward pass I simply iterate the LSTMCell over sequence length and the num_layers as follow:
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]))
Where input is of size (batch_size, input_size), prev_hiddens[i] is size of (batch_size, hidden_size), prev_cells[i] is of size (batch_size, hidden_size).
In the backward_hook I print the size of the tensors that are input to this function:
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()))
As the results, for the first time backward_hook is called for example:
[A] For grad_output I get 2 tensors among which the second tensor is None. This is understandable because in the backward phase we have a gradient of internal states (c) and gradient of output (h). The last iteration in time dimension has no future hidden so its gradient is None.
[B] For grad_input I get 5 tensors (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])
My questions are:
(1) Is my understanding from [A] correct?
(2) How do I interpret the 5 tensors from the grad_input tuple? I thought there should have only 3 since there are only 3 inputs to the LSTMCell forward()?
Thanks

Answer 1

Your understanding of grad_input and grad_output is wrong. I am trying to explain it with a simpler example.
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()
The output is:
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])
Explanation
grad_output: Gradient of the loss w.r.t. the layer output, Y_pred.
grad_input: Gradients of the loss w.r.t the layer inputs. For Linear layer, the inputs are the input tensor and the weight and the bias.
So, in the output you see:
grad_input.size: torch.Size([8, 20]) # for the `bias`
None # for the `input`
grad_input.size: torch.Size([10, 20]) # for the `weight`
The Linear layer in PyTorch uses a LinearFunction which is as follows.
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
For LSTM, there are four sets of weight parameters.
weight_ih_l0
weight_hh_l0
bias_ih_l0
bias_hh_l0
So, in your case, the grad_input would be a tuple of 5 tensors. And as you mentioned, the grad_output is two tensors.