def test_cudnn_rnn_speed(self):
from time import time
iters = 100
h0 = np.random.rand(1, 128, 256).astype(np.float32)
input = np.random.rand(128, 128, 128).astype(np.float32)
dev = torch.device('cuda:0')
t_rnn = tnn.RNN(128, 256, nonlinearity='relu').to(dev)
t_optim = torch.optim.SGD(t_rnn.parameters(), lr=1e-3, momentum=0.9)
t_input = torch.from_numpy(input).to(dev)
t_h0 = torch.from_numpy(h0).to(dev)
start_time = time()
for i in range(iters):
t_optim.zero_grad()
t_output, th = t_rnn(t_input, t_h0)
t_loss = (t_output**2).sum() + (th**2).sum()
t_loss.backward()
t_optim.step()
print('torch time = ', time() - start_time)
j_rnn = nn.RNN(128, 256, nonlinearity='relu')
j_rnn.load_state_dict(t_rnn.state_dict())
j_optim = nn.SGD(j_rnn.parameters(), lr=1e-3, momentum=0.9)
j_input, j_h0 = jt.array(input), jt.array(h0)
start_time = time()
for i in range(iters):
j_output, jh = j_rnn(j_input, j_h0)
j_loss = (j_output**2).sum() + (jh**2).sum()
j_optim.step(j_loss)
jt.sync_all(True)
print('jittor Cudnn time = ', time() - start_time)
jt_cudnn, jt.cudnn = jt.cudnn, None
j_rnn = nn.RNN(128, 256, nonlinearity='relu')
j_rnn.load_state_dict(t_rnn.state_dict())
j_optim = nn.SGD(j_rnn.parameters(), lr=1e-3, momentum=0.9)
start_time = time()
for i in range(iters):
j_output, jh = j_rnn(j_input, j_h0)
j_loss = (j_output**2).sum() + (jh**2).sum()
j_optim.step(j_loss)
jt.sync_all(True)
print('jittor native time = ', time() - start_time)
jt.cudnn = jt_cudnn
def test_cudnn_rnn_train(self):
dev = torch.device('cuda:0')
t_rnn = tnn.RNN(32, 64, nonlinearity='relu').to(dev)
t_optim = torch.optim.SGD(t_rnn.parameters(), lr=1e-3, momentum=0.9)
j_rnn = nn.RNN(32, 64, nonlinearity='relu')
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(1, 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 = t_rnn(
torch.from_numpy(input).to(dev),
torch.from_numpy(h0).to(dev))
t_loss = (t_output**2).sum() + (th**2).sum()
t_loss.backward()
t_optim.step()
j_input, jh = jt.array(input), jt.array(h0)
j_output, jh = j_rnn(j_input, jh)
j_loss = (j_output**2).sum() + (jh**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_no_bias_rnn(self):
h0 = np.random.rand(4, 4, 200).astype(np.float32)
input = np.random.rand(5, 4, 100).astype(np.float32)
t_rnn = tnn.RNN(100, 200, num_layers=2, bidirectional=True, bias=False)
j_rnn = nn.RNN(100, 200, num_layers=2, bidirectional=True, bias=False)
check_equal_1(t_rnn, j_rnn, input, h0)
def test_multilayer_rnn(self):
h0 = np.random.rand(4, 4, 200).astype(np.float32)
input = np.random.rand(5, 4, 100).astype(np.float32)
t_rnn = tnn.RNN(100, 200, num_layers=4)
j_rnn = nn.RNN(100, 200, num_layers=4)
check_equal_1(t_rnn, j_rnn, input, h0)
def test_bidirectional_rnn(self):
h0 = np.random.rand(2, 1, 200).astype(np.float32)
input = np.random.rand(5, 1, 100).astype(np.float32)
t_rnn = tnn.RNN(100, 200, bidirectional=True)
j_rnn = nn.RNN(100, 200, bidirectional=True)
check_equal_1(t_rnn, j_rnn, input, h0)
def test_basic_rnn(self):
h0 = np.random.rand(1, 24, 200).astype(np.float32)
input = np.random.rand(32, 24, 100).astype(np.float32)
t_rnn = tnn.RNN(100, 200)
j_rnn = nn.RNN(100, 200)
check_equal_1(t_rnn, j_rnn, input, h0)
def test_multilayer_cudnn_rnn(self):
dev = torch.device('cuda:0')
t_rnn = tnn.RNN(100, 200, num_layers=4, nonlinearity='tanh').to(dev)
j_rnn = nn.RNN(100, 200, num_layers=4, nonlinearity='tanh')
h0 = np.random.rand(4, 8, 200).astype(np.float32)
input = np.random.rand(5, 8, 100).astype(np.float32)
check_equal_1(t_rnn, j_rnn, input, h0, dev)
def test_basic_cudnn_rnn(self):
dev = torch.device('cuda:0')
t_rnn = tnn.RNN(100, 200, nonlinearity='relu').to(dev)
j_rnn = nn.RNN(100, 200, nonlinearity='relu')
h0 = np.random.rand(1, 24, 200).astype(np.float32)
input = np.random.rand(32, 24, 100).astype(np.float32)
check_equal_1(t_rnn, j_rnn, input, h0, dev)
def test_no_bias_cudnn_rnn(self):
dev = torch.device('cuda:0')
t_rnn = tnn.RNN(100, 200, bidirectional=True, bias=False, nonlinearity='tanh').to(dev)
j_rnn = nn.RNN(100, 200, bidirectional=True, bias=False, nonlinearity='tanh')
h0 = np.random.rand(2, 8, 200).astype(np.float32)
input = np.random.rand(5, 8, 100).astype(np.float32)
check_equal_1(t_rnn, j_rnn, input, h0, dev)
def test_dropout_rnn(self):
h0 = np.random.rand(2, 4, 200).astype(np.float32)
input = np.random.rand(5, 4, 100).astype(np.float32)
t_rnn = tnn.RNN(100, 200, num_layers=2, dropout=0.5, bias=False)
j_rnn = nn.RNN(100, 200, num_layers=2, dropout=0.5, bias=False)
t_rnn.eval()
j_rnn.eval()
check_equal_1(t_rnn, j_rnn, input, h0)
def test_rnn_default_hx(self):
input = np.random.rand(32, 24, 12).astype(np.float32)
h0 = np.zeros((1, 24, 24)).astype(np.float32)
t_rnn = tnn.RNN(12, 24)
j_rnn = nn.RNN(12, 24)
j_rnn.load_state_dict(t_rnn.state_dict())
t_output, th = t_rnn(torch.from_numpy(input))
j_output, jh = 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)
def test_rnn(self):
h0 = np.random.rand(1, 24, 200).astype(np.float32)
input = np.random.rand(32, 24, 100).astype(np.float32)
t_rnn = tnn.RNN(100, 200)
j_rnn = nn.RNN(100, 200)
check_equal_1(t_rnn, j_rnn, input, h0)
h0 = np.random.rand(4, 4, 200).astype(np.float32)
input = np.random.rand(5, 4, 100).astype(np.float32)
t_rnn = tnn.RNN(100, 200, num_layers=4)
j_rnn = nn.RNN(100, 200, num_layers=4)
check_equal_1(t_rnn, j_rnn, input, h0)
h0 = np.random.rand(2, 1, 200).astype(np.float32)
input = np.random.rand(5, 1, 100).astype(np.float32)
t_rnn = tnn.RNN(100, 200, bidirectional=True)
j_rnn = nn.RNN(100, 200, bidirectional=True)
check_equal_1(t_rnn, j_rnn, input, h0)
h0 = np.random.rand(4, 4, 200).astype(np.float32)
input = np.random.rand(5, 4, 100).astype(np.float32)
t_rnn = tnn.RNN(100, 200, num_layers=2, bidirectional=True, bias=False)
j_rnn = nn.RNN(100, 200, num_layers=2, bidirectional=True, bias=False)
check_equal_1(t_rnn, j_rnn, input, h0)
h0 = np.random.rand(2, 4, 200).astype(np.float32)
input = np.random.rand(5, 4, 100).astype(np.float32)
t_rnn = tnn.RNN(100, 200, num_layers=2, dropout=0.5, bias=False)
j_rnn = nn.RNN(100, 200, num_layers=2, dropout=0.5, bias=False)
t_rnn.eval()
j_rnn.eval()
check_equal_1(t_rnn, j_rnn, input, h0)
def test_cudnn_rnn(self):
dev = torch.device('cuda:0')
t_rnn = tnn.RNN(100, 200, nonlinearity='relu').to(dev)
j_rnn = nn.RNN(100, 200, nonlinearity='relu')
j_rnn.train()
j_rnn.load_state_dict(t_rnn.state_dict())
h0 = np.random.rand(1, 24, 200).astype(np.float32)
input = np.random.rand(32, 24, 100).astype(np.float32)
t_output, th = t_rnn(torch.from_numpy(input).to(dev),
torch.from_numpy(h0).to(dev))
j_output, jh = j_rnn(jt.array(input), jt.array(h0))
np.testing.assert_allclose(j_output.data, t_output.detach().cpu().numpy())
np.testing.assert_allclose(jh.data, th.detach().cpu().numpy())
