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. 总结
- warp中labels的size应是N*S或 sum(target_lens)
- torch.nn中的labels的size应为(N,S)或sum(target_lens)
- 区别似乎主要在log_softmax()
- 尽量还是先用warpctc