LSTM stepwise backward

lstm_backward.py

lstm = torch.lstm_cell
grad_history = []

def store_grads(name, timestep):
    def hook(grad):
        print(name, grad, timestep)
        grad_history[timestep][name] = grad
    return hook

def lstm_stepwise(input, h0, w_ih, w_hh, b_ih, b_hh):
    h = h0
    T = input.size(1)
    for t in range(T):
        grad_history.append({})
        
        w_ih.register_hook(store_grads('w_ih', t))
        w_hh.register_hook(store_grads('w_hh', t))
        b_ih.register_hook(store_grads('b_ih', t))
        b_hh.register_hook(store_grads('b_hh', t))
        
        h = lstm(input[:, t, :], h, w_ih, w_hh, b_ih, b_hh)
    return h

out = lstm_stepwise(data, h0, w_ih, w_hh, b_ih, b_hh)
loss = criterion(out[0], y)
loss.backward()