def test_model():
freq_dim = 120
vocab_size = 10
np.random.seed(1337)
torch.manual_seed(1337)
conf = shared.model_config
rnn_dim = conf['encoder']['rnn']['dim']
conf["decoder"] = {"embedding_dim": rnn_dim, "layers": 2}
model = Seq2Seq(freq_dim, vocab_size + 1, conf)
batch = shared.gen_fake_data(freq_dim, vocab_size)
batch_size = len(batch[0])
out = model(batch)
loss = model.loss(batch)
assert out.size()[0] == batch_size
assert out.size()[2] == vocab_size
assert len(out.size()) == 3
x, y = model.collate(*batch)
x_enc = model.encode(x)
state = None
out_s = []
for t in range(y.size()[1] - 1):
ox, state = model.decode_step(x_enc, y[:, t:t + 1], state=state)
out_s.append(ox)
out_s = torch.stack(out_s, dim=1)
assert out.size() == out_s.size()
assert np.allclose(out_s.data.numpy(),
out.data.numpy(),
rtol=1e-5,
atol=1e-7)
def test_ctc_model():
freq_dim = 40
output_dim = 10
batch = shared.gen_fake_data(freq_dim, output_dim)
batch_size = len(batch[0])
model = CTC(freq_dim, output_dim, shared.model_config)
out = model(batch)
assert out.size()[0] == batch_size
assert out.size()[2] == output_dim
assert len(out.size()) == 3
loss = model.loss(batch)
preds = model.infer(batch)
assert len(preds) == batch_size
def test_ctc_model():
freq_dim = 40
vocab_size = 10
batch = shared.gen_fake_data(freq_dim, vocab_size)
batch_size = len(batch[0])
model = CTC(freq_dim, vocab_size, shared.model_config)
out = model(batch)
assert out.size()[0] == batch_size
# CTC model adds the blank token to the vocab
assert out.size()[2] == (vocab_size + 1)
assert len(out.size()) == 3
loss = model.loss(batch)
preds = model.infer(batch)
assert len(preds) == batch_size