def bench():
set_gpu()
set_allow_cuda(True)
T, B, F = 100, 64, 300
n_trials = 10
fake_input = cuda_move(torch.zeros(T, B).long())
print("JitRNN ", end='')
rnn = RNNLayer(100, 100)
model = cuda_move(DiscreteRNN(F, F, 100, rnn=rnn))
model(fake_input)
foo = lambda: model(fake_input)
timeit(foo, n_trials)
print("JitLMN ", end='')
rnn = LMNLayer(100, 100, 100)
model = cuda_move(DiscreteRNN(F, F, 100, rnn=rnn))
model(fake_input)
foo = lambda: model(fake_input)
timeit(foo, n_trials)
print("JitSlowLMN ", end='')
rnn = LMNLayer(100, 100, 100)
model = cuda_move(DiscreteRNN(F, F, 100, rnn=rnn))
model(fake_input)
foo = lambda: model(fake_input)
timeit(foo, n_trials)
print("SlowLMN ", end='')
rnn = SlowLMNLayer(100, 100, 100)
model = cuda_move(SlowDiscreteRNN(F, F, 100, rnn=rnn))
model(fake_input)
foo = lambda: model(fake_input)
timeit(foo, n_trials)
print("SlowLSTM ", end='')
rnn = SlowLSTM(input_size=100, hidden_size=100, num_layers=1)
model = cuda_move(SlowDiscreteRNN(F, F, 100, rnn=rnn))
model(fake_input)
foo = lambda: model(fake_input)
timeit(foo, n_trials)
from cannon.tasks import CopyTask, AdditionTask
from cannon.utils import set_gpu, set_allow_cuda
import torch
set_gpu()
def test_copy():
set_allow_cuda(False)
S, T, K = 10, 5, 2
batch_size = 3
data = CopyTask(S, T, K)
x, y = data.get_batch(batch_size)
# marker at correct position
assert (x[T + S - 1] == K + 1).all()
mask = torch.cat([y[:T + S], x[:S]], dim=0)
out = torch.zeros((T + 2 * S, batch_size, K + 2))
# optimal solution has zero error (logits should be +inf).
for t in range(2 * S + T):
for b in range(batch_size):
out[t, b, mask[t, b]] = 10**9
loss = data.loss_score(out, y)
assert loss == 0.0
def test_addition():
seq_len = 10
batch_size = 4