torch.nn.CTCLoss 与warpctc_pytoch.CTCLoss

1.torch.nn.CTCLoss

import torch
from torch.nn import CTCLoss

torch.backends.cudnn.benchmark = True
T = 50      # Input sequence length
C = 20      # Number of classes (including blank)
N = 16      # Batch size
S = 30     # Target sequence length of longest target in batch
S_min = 10
torch.manual_seed(1234)

input = torch.randn(T, N, C).log_softmax(2).detach().requires_grad_()
target = torch.randint(low=1, high=C, size=(N,S), dtype=torch.long)
input_lengths = torch.full(size=(N,), fill_value=T, dtype=torch.long)
target_lengths = torch.randint(low=S_min, high=S, size=(N,), dtype=torch.long)

critenzer = CTCLoss()
loss = critenzer(input,target,input_lengths,target_lengths)
loss.backward()
loss

output:

tensor(10.5236, grad_fn=< MeanBackward0>)

2.warpctc_pytorch CTCLoss

源码见github: warpctc-pytorch

from warpctc_pytorch import CTCLoss as ctc

probs = torch.FloatTensor([[[0.1, 0.6, 0.1, 0.1, 0.1], [0.1, 0.1, 0.6, 0.1, 0.1]]]).transpose(0, 1).contiguous()
labels = torch.IntTensor([1, 2])
label_sizes = torch.IntTensor([2])
probs_sizes = torch.IntTensor([2])
probs.requires_grad_(True)  # tells autograd to compute gradients for probs

output:

tensor([2.4629], grad_fn=<_CTCBackward>)

    probs: Tensor of (seqLength x batch x outputDim) containing output from network
    
    labels: 1 dimensional Tensor containing all the targets of the batch in one sequence
    
    probs_lens: Tensor of size (batch) containing size of each output sequence from the network
    
    label_lens: Tensor of (batch) containing label length of each example

3. 总结

1. warp中labels的size应是N*S或 sum(target_lens)
2. torch.nn中的labels的size应为(N,S)或sum(target_lens)
3. 区别似乎主要在log_softmax()
4. 尽量还是先用warpctc