Test the bidirectional LSTM type

bilstm_type.py

"""A simple script to test the biLSTM type that pytorch uses.
The gradients are computed only w.r.t the output of one single direction,
so gradient of the reverse direction in layer 1 should be zero if type1.

In my tests, it's type2
"""

import torch
from torch import nn


bilstm = nn.LSTM(10, 100, 2, bidirectional=True)
fake_input = torch.Tensor(10, 10, 10)

output, _ = bilstm(fake_input)
output[..., :100].sum().backward()

print('==========BiLSTM on CPU==========')
for name, param in bilstm.named_parameters():
    print('{} grad: {}'.format(name, param.grad.mean().item()))

bilstm.zero_grad()
bilstm.cuda()
output, _ = bilstm(fake_input.cuda())
output[..., :100].sum().backward()

print('==========BiLSTM on GPU==========')
for name, param in bilstm.named_parameters():
    print('{} grad: {}'.format(name, param.grad.mean().item()))