Add rnnt loss

[pkufool]

• Add mutual information
• Add rnnt loss
• Add unit tests
• Update documents and fix code style

[danpovey]

Oh, this is something I forgot to test...
More ideally this would be #if 0, that path is slightly more optimized, but I have not tested it.

[pkufool]

Ok, I will test it.

[pkufool]

I think this rnnt loss runs fine now, I will update the experiment results in Icefall (k2-fsa/icefall#170).

[danpovey]

Here, instead of the torch.stack expression, it would be more efficient to do
torch.arange(0, T + 1, device=px_am.device).reshape(1, T + 1).expand(B, T + 1)
which is a zero-copy expression.
Also, below, instead of doing mask = torch.where(mask < boundary[:, 3].reshape(B, 1), True, False)
you can just use mask = mask < boundary[:, 3].reshape(B, 1) , which is already a boolean expression.

[csukuangfj]

Suggested change

	ans_grad = ans_grad.reshape((B, 1, 1)) # (B, 1, 1)
	ans_grad = ans_grad.reshape(B, 1, 1) # (B, 1, 1)

[csukuangfj]

Suggested change

	assert px.stride()[-1] == 1
	assert px.stride(-1) == 1

[csukuangfj]

I would recommend

assert px.is_contiguous()
assert py.is_contiguous()

[csukuangfj]

If return_grad is False, is it possible to not return px_grad and py_grad to save memory?

[danpovey]

... that is a valid point. In some circumstances we might not even need the backprop. I'm pretty sure there was a time when it allowed to delay the backward part of the computation, it must have been simplified since then.
The memory requirements are not very huge because it's only one element per (s,t) position, not C elements, but it's still nonzero.

[csukuangfj]

Suggested change

	assert a.shape[-1] == b.shape[-1] # last last dim be K
	assert a.shape[-1] == b.shape[-1] # The last dim must be equal

[csukuangfj]

Suggested change

	def _inner(a: Tensor, b: Tensor) -> Tensor:
	def _inner_product(a: Tensor, b: Tensor) -> Tensor:

[csukuangfj]

I would recommend using torch::optional<torch::Tensor>.
You don't need to create a tensor of shape (0, 0). Just leave it unset.

Also, you can pass a None from Python.

[csukuangfj]

Suggested change

	px_grad, py_grad = px_grad.reshape(1, B, -1), py_grad.reshape(1, B, -1)
	px_grad = px_grad.reshape(1, B, -1)
	py_grad = py_grad.reshape(1, B, -1)

is more readable.

[csukuangfj]

Suggested change

	px_cat, py_cat = px_cat.clamp(min=-1.0e+38), py_cat.clamp(min=-1.0e+38)
	px_cat, py_cat = px_cat.clamp(min=torch.finfo(torch.float32).min), py_cat.clamp(min=-1.0e+38)

Instead of doing

a, b = op(a), op(b)

I would recommend doing only one thing in a single line

a = op(a)
b = op(b)

[csukuangfj]

I would recommend separating these conjunctions out.
If one of the assertions fails, you don't know exactly which one failed.

[csukuangfj]

What is the magic number 1.0e-37?

[danpovey]

It's close to the lowest representable single precision float. I was concerned making this too high, like 1e-20, might be the cause of instability.

[csukuangfj]

There is

torch.finfo(torch.float32).tiny

which is 1.17549e-38, Is this the value you want?

[danpovey]

Yes I suppose so. Could even use the dtype of lm_probs instead of torch.float32.

[pkufool]

Merging now, as it will be more convenient for me to integrate the modified version.

[csukuangfj]

p(x_s | x_0..x_{s-1}, y_0..y_{s-1})

--> (the last one: s-1 --> t-1)

p(x_s | x_0..x_{s-1}, y_0..y_{t-1})

[csukuangfj]

typo: directy -> directly

[csukuangfj]

https://github.com/pkufool/k2/blob/5dbf4081a99938cc9e3052de0bbb069346128725/k2/python/k2/mutual_information.py#L132
says:

    with 0 <= s_begin <= s_end < S and 0 <= t_begin <= t_end < T

But this line is using

            assert 0 <= s_begin <= s_end <= S

Which one is correct?

[danpovey]

should be <=.

[csukuangfj]

https://github.com/pkufool/k2/blob/5dbf4081a99938cc9e3052de0bbb069346128725/k2/python/k2/mutual_information.py#L114
says

    Note: we don't require px and py to be contiguous, but the
    code assumes for optimization purposes that the T axis has
    stride 1.

Buth here it is using

    assert px.is_contiguous()
    assert py.is_contiguous()

Which one is correct?

[danpovey]

We could probably reduce the assertions to: `assert px.stride(-1) == 1 and py.stride(-1) == 1

[csukuangfj]

I find that the implementation is using tensor's accessor to access the underlying data, which handles the strides automagically. So I think it does not pose any constraints on the layout of px and py.

We don't need to check the stride of px and py here, I think.

[EDITED]: At least it is true for the CPU implementation.

[danpovey]

It may not actually be required (although I would have to check the GPU implementation), but when the code was designed, I assumed it had stride 1, and it would likely be super slow if not.

[csukuangfj]

torch.logsumexp() is numerically stable, so there is no need to subtract the max here.

See https://pytorch.org/docs/stable/generated/torch.logsumexp.html

Returns the log of summed exponentials of each row of the input tensor in the given dimension dim. The computation is numerically stabilized.

and see its internal implementation
https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/ReduceOps.cpp#L1260

static Tensor& logsumexp_out_impl(Tensor& result, const Tensor& self, IntArrayRef dims, bool keepdim) {
  // can't take max of empty tensor
  if (self.numel() != 0) {
    auto maxes = at::amax(self, dims, true);
    auto maxes_squeezed = (keepdim ? maxes : squeeze_multiple(maxes, dims));
    maxes_squeezed.masked_fill_(maxes_squeezed.abs() == INFINITY, 0);
    at::sum_out(result, (self - maxes).exp_(), dims, keepdim);
    result.log_().add_(maxes_squeezed);
  } else {

[pkufool]

Yes, Dan had a comment on this line before, I forgot to delete it. Will fix the comments you gave these days. Thanks!

[csukuangfj]

I see. It has been fixed in the master.

[csukuangfj]

You are using the boundary that was defined when device is cpu.

The code should work if users disable the tests for CPU.