def test_lstm(self):
h0 = np.random.rand(1, 2, 20).astype(np.float32)
c0 = np.random.rand(1, 2, 20).astype(np.float32)
input = np.random.rand(5, 2, 10).astype(np.float32)
t_rnn = tnn.LSTM(10, 20)
j_rnn = nn.LSTM(10, 20)
check_equal(t_rnn, j_rnn, input, h0, c0)
proj_size = 13
h0 = np.random.rand(1, 2, proj_size).astype(np.float32)
c0 = np.random.rand(1, 2, 20).astype(np.float32)
input = np.random.rand(5, 2, 10).astype(np.float32)
t_rnn = tnn.LSTM(10, 20, proj_size=proj_size)
j_rnn = nn.LSTM(10, 20, proj_size=proj_size)
check_equal(t_rnn, j_rnn, input, h0, c0)
h0 = np.random.rand(2, 4, 20).astype(np.float32)
c0 = np.random.rand(2, 4, 20).astype(np.float32)
input = np.random.rand(5, 4, 10).astype(np.float32)
t_rnn = tnn.LSTM(10, 20, num_layers=2)
j_rnn = nn.LSTM(10, 20, num_layers=2)
check_equal(t_rnn, j_rnn, input, h0, c0)
h0 = np.random.rand(2, 4, proj_size).astype(np.float32)
c0 = np.random.rand(2, 4, 20).astype(np.float32)
input = np.random.rand(5, 4, 10).astype(np.float32)
t_rnn = tnn.LSTM(10, 20, num_layers=2, proj_size=proj_size)
j_rnn = nn.LSTM(10, 20, num_layers=2, proj_size=proj_size)
check_equal(t_rnn, j_rnn, input, h0, c0)
def test_twobilinear_lstm(self):
x = jt.rand(5, 4, 10)
rnn1 = nn.LSTM(10, 20, bidirectional=True)
out1, _ = rnn1(x)
rnn2 = nn.LSTM(40, 20, bidirectional=True)
out2, _ = rnn2(out1)
target = jt.zeros_like(out2)
loss = nn.mse_loss(out2, target)
from jittor import optim
optimizer = optim.RMSprop(rnn1.parameters())
optimizer.step(loss)
def test_multilayer_bidirectional_cudnn_lstm_train(self):
dev = torch.device('cuda:0')
t_rnn = tnn.LSTM(32, 64, num_layers=4, bidirectional=True).to(dev)
t_optim = torch.optim.SGD(t_rnn.parameters(), lr=1e-3, momentum=0.9)
j_rnn = nn.LSTM(32, 64, num_layers=4, bidirectional=True)
j_rnn.load_state_dict(t_rnn.state_dict())
j_optim = nn.SGD(j_rnn.parameters(), lr=1e-3, momentum=0.9)
h0 = np.random.rand(8, 4, 64).astype(np.float32)
c0 = np.random.rand(8, 4, 64).astype(np.float32)
input = np.random.rand(12, 4, 32).astype(np.float32)
for _ in range(10):
t_optim.zero_grad()
t_output, (th, tc) = t_rnn(
torch.from_numpy(input).to(dev),
(torch.from_numpy(h0).to(dev), torch.from_numpy(c0).to(dev)))
t_loss = (t_output**2).sum() + (th**2).sum() + (tc**2).sum()
t_loss.backward()
t_optim.step()
j_input, jh0, jc0 = jt.array(input), jt.array(h0), jt.array(c0)
j_output, (jh, jc) = j_rnn(j_input, (jh0, jc0))
j_loss = (j_output**2).sum() + (jh**2).sum() + (jc**2).sum()
j_optim.step(j_loss)
np.testing.assert_allclose(t_loss.item(), j_loss.item(), rtol=1e-4)
np.testing.assert_allclose(t_rnn.bias_hh_l0.detach().cpu().numpy(),
j_rnn.bias_hh_l0.data,
atol=1e-4,
rtol=1e-4)
def test_bidirectional_lstm(self):
h0 = np.random.rand(2, 1, 200).astype(np.float32)
c0 = np.random.rand(2, 1, 200).astype(np.float32)
input = np.random.rand(5, 1, 100).astype(np.float32)
t_rnn = tnn.LSTM(100, 200, bidirectional=True)
j_rnn = nn.LSTM(100, 200, bidirectional=True)
check_equal_2(t_rnn, j_rnn, input, h0, c0)
def test_multilayer_lstm(self):
h0 = np.random.rand(4, 4, 200).astype(np.float32)
c0 = np.random.rand(4, 4, 200).astype(np.float32)
input = np.random.rand(5, 4, 100).astype(np.float32)
t_rnn = tnn.LSTM(100, 200, num_layers=4)
j_rnn = nn.LSTM(100, 200, num_layers=4)
check_equal_2(t_rnn, j_rnn, input, h0, c0)
def test_projection_lstm(self):
proj_size = 13
h0 = np.random.rand(1, 24, proj_size).astype(np.float32)
c0 = np.random.rand(1, 24, 200).astype(np.float32)
input = np.random.rand(32, 24, 100).astype(np.float32)
t_rnn = tnn.LSTM(100, 200, proj_size=proj_size)
j_rnn = nn.LSTM(100, 200, proj_size=proj_size)
check_equal_2(t_rnn, j_rnn, input, h0, c0)
def test_basic_lstm(self):
h0 = np.random.rand(1, 24, 200).astype(np.float32)
c0 = np.random.rand(1, 24, 200).astype(np.float32)
input = np.random.rand(32, 24, 100).astype(np.float32)
t_rnn = tnn.LSTM(100, 200)
j_rnn = nn.LSTM(100, 200)
check_equal_2(t_rnn, j_rnn, input, h0, c0)
def test_multilayer_bidirectional_projection_lstm(self):
h0 = np.random.rand(4, 4, 200).astype(np.float32)
c0 = np.random.rand(4, 4, 200).astype(np.float32)
input = np.random.rand(5, 4, 100).astype(np.float32)
t_rnn = tnn.LSTM(100, 200, num_layers=2, bidirectional=True, bias=False)
j_rnn = nn.LSTM(100, 200, num_layers=2, bidirectional=True, bias=False)
check_equal_2(t_rnn, j_rnn, input, h0, c0)
def test_basic_lstm_rnn(self):
dev = torch.device('cuda:0')
t_rnn = tnn.LSTM(100, 200).to(dev)
j_rnn = nn.LSTM(100, 200)
h0 = np.random.rand(1, 24, 200).astype(np.float32)
c0 = np.random.rand(1, 24, 200).astype(np.float32)
input = np.random.rand(32, 24, 100).astype(np.float32)
check_equal_2(t_rnn, j_rnn, input, h0, c0, dev)
def test_bidirectional_projection_lstm(self):
proj_size = 10
h0 = np.random.rand(2, 4, proj_size).astype(np.float32)
c0 = np.random.rand(2, 4, 200).astype(np.float32)
input = np.random.rand(5, 4, 100).astype(np.float32)
t_rnn = tnn.LSTM(100, 200, bidirectional=True, proj_size=proj_size)
j_rnn = nn.LSTM(100, 200, bidirectional=True, proj_size=proj_size)
check_equal_2(t_rnn, j_rnn, input, h0, c0)
def test_multilayer_projection_lstm(self):
proj_size = 8
h0 = np.random.rand(2, 4, proj_size).astype(np.float32)
c0 = np.random.rand(2, 4, 20).astype(np.float32)
input = np.random.rand(5, 4, 10).astype(np.float32)
t_rnn = tnn.LSTM(10, 20, num_layers=2, proj_size=proj_size)
j_rnn = nn.LSTM(10, 20, num_layers=2, proj_size=proj_size)
check_equal_2(t_rnn, j_rnn, input, h0, c0)
def test_dropout_lstm(self):
h0 = np.random.rand(2, 4, 200).astype(np.float32)
c0 = np.random.rand(2, 4, 200).astype(np.float32)
input = np.random.rand(5, 4, 100).astype(np.float32)
t_rnn = tnn.LSTM(100, 200, num_layers=2, dropout=0.5, bias=False)
j_rnn = nn.LSTM(100, 200, num_layers=2, dropout=0.5, bias=False)
t_rnn.eval()
j_rnn.eval()
check_equal_2(t_rnn, j_rnn, input, h0, c0)
def test_lstm_default_hx(self):
input = np.random.rand(32, 24, 10).astype(np.float32)
t_rnn = tnn.LSTM(10, 20, num_layers=2, bidirectional=True)
j_rnn = nn.LSTM(10, 20, num_layers=2, bidirectional=True)
j_rnn.load_state_dict(t_rnn.state_dict())
t_output, (th, tc) = t_rnn(torch.from_numpy(input))
j_output, (jh, jc) = j_rnn(jt.array(input))
np.testing.assert_allclose(t_output.detach().cpu().numpy(), j_output.data, rtol=1e-03, atol=1e-06)
np.testing.assert_allclose(th.detach().cpu().numpy(), jh.data, rtol=1e-03, atol=1e-06)
np.testing.assert_allclose(tc.detach().cpu().numpy(), jc.data, rtol=1e-03, atol=1e-06)
def test_lstm(self):
h0 = np.random.rand(1, 24, 200).astype(np.float32)
c0 = np.random.rand(1, 24, 200).astype(np.float32)
input = np.random.rand(32, 24, 100).astype(np.float32)
t_rnn = tnn.LSTM(100, 200)
j_rnn = nn.LSTM(100, 200)
check_equal_2(t_rnn, j_rnn, input, h0, c0)
proj_size = 13
h0 = np.random.rand(1, 24, proj_size).astype(np.float32)
c0 = np.random.rand(1, 24, 200).astype(np.float32)
input = np.random.rand(32, 24, 100).astype(np.float32)
t_rnn = tnn.LSTM(100, 200, proj_size=proj_size)
j_rnn = nn.LSTM(100, 200, proj_size=proj_size)
check_equal_2(t_rnn, j_rnn, input, h0, c0)
h0 = np.random.rand(4, 4, 200).astype(np.float32)
c0 = np.random.rand(4, 4, 200).astype(np.float32)
input = np.random.rand(5, 4, 100).astype(np.float32)
t_rnn = tnn.LSTM(100, 200, num_layers=4)
j_rnn = nn.LSTM(100, 200, num_layers=4)
check_equal_2(t_rnn, j_rnn, input, h0, c0)
h0 = np.random.rand(2, 4, proj_size).astype(np.float32)
c0 = np.random.rand(2, 4, 20).astype(np.float32)
input = np.random.rand(5, 4, 10).astype(np.float32)
t_rnn = tnn.LSTM(10, 20, num_layers=2, proj_size=proj_size)
j_rnn = nn.LSTM(10, 20, num_layers=2, proj_size=proj_size)
check_equal_2(t_rnn, j_rnn, input, h0, c0)
h0 = np.random.rand(2, 1, 200).astype(np.float32)
c0 = np.random.rand(2, 1, 200).astype(np.float32)
input = np.random.rand(5, 1, 100).astype(np.float32)
t_rnn = tnn.LSTM(100, 200, bidirectional=True)
j_rnn = nn.LSTM(100, 200, bidirectional=True)
check_equal_2(t_rnn, j_rnn, input, h0, c0)
proj_size = 13
h0 = np.random.rand(2, 4, proj_size).astype(np.float32)
c0 = np.random.rand(2, 4, 200).astype(np.float32)
input = np.random.rand(5, 4, 100).astype(np.float32)
t_rnn = tnn.LSTM(100, 200, bidirectional=True, proj_size=proj_size)
j_rnn = nn.LSTM(100, 200, bidirectional=True, proj_size=proj_size)
check_equal_2(t_rnn, j_rnn, input, h0, c0)
h0 = np.random.rand(4, 4, 200).astype(np.float32)
c0 = np.random.rand(4, 4, 200).astype(np.float32)
input = np.random.rand(5, 4, 100).astype(np.float32)
t_rnn = tnn.LSTM(100, 200, num_layers=2, bidirectional=True, bias=False)
j_rnn = nn.LSTM(100, 200, num_layers=2, bidirectional=True, bias=False)
check_equal_2(t_rnn, j_rnn, input, h0, c0)
h0 = np.random.rand(2, 4, 200).astype(np.float32)
c0 = np.random.rand(2, 4, 200).astype(np.float32)
input = np.random.rand(5, 4, 100).astype(np.float32)
t_rnn = tnn.LSTM(100, 200, num_layers=2, dropout=0.5, bias=False)
j_rnn = nn.LSTM(100, 200, num_layers=2, dropout=0.5, bias=False)
t_rnn.eval()
j_rnn.eval()
check_equal_2(t_rnn, j_rnn, input, h0, c0)
