mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-27 08:58:13 +00:00
Brandon Thomas
71 lines
2.6 KiB
Python
71 lines
2.6 KiB
Python
# ! /usr/bin/python
|
|
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import torch
|
|
from torch import nn
|
|
|
|
from nemo.core.classes import Serialization, Typing, typecheck
|
|
from nemo.core.neural_types import LabelsType, LengthsType, LogprobsType, LossType, NeuralType
|
|
|
|
__all__ = ['CTCLoss']
|
|
|
|
|
|
class CTCLoss(nn.CTCLoss, Serialization, Typing):
|
|
@property
|
|
def input_types(self):
|
|
"""Input types definitions for CTCLoss.
|
|
"""
|
|
return {
|
|
"log_probs": NeuralType(('B', 'T', 'D'), LogprobsType()),
|
|
"targets": NeuralType(('B', 'T'), LabelsType()),
|
|
"input_lengths": NeuralType(tuple('B'), LengthsType()),
|
|
"target_lengths": NeuralType(tuple('B'), LengthsType()),
|
|
}
|
|
|
|
@property
|
|
def output_types(self):
|
|
"""Output types definitions for CTCLoss.
|
|
loss:
|
|
NeuralType(None)
|
|
"""
|
|
return {"loss": NeuralType(elements_type=LossType())}
|
|
|
|
def __init__(self, num_classes, zero_infinity=False, reduction='mean_batch'):
|
|
self._blank = num_classes
|
|
# Don't forget to properly call base constructor
|
|
if reduction == 'mean_batch':
|
|
ctc_reduction = 'none'
|
|
self._apply_batch_mean = True
|
|
elif reduction in ['sum', 'mean', 'none']:
|
|
ctc_reduction = reduction
|
|
self._apply_batch_mean = False
|
|
super().__init__(blank=self._blank, reduction=ctc_reduction, zero_infinity=zero_infinity)
|
|
|
|
@typecheck()
|
|
def forward(self, log_probs, targets, input_lengths, target_lengths):
|
|
# override forward implementation
|
|
# custom logic, if necessary
|
|
input_lengths = input_lengths.long()
|
|
target_lengths = target_lengths.long()
|
|
targets = targets.long()
|
|
# here we transpose because we expect [B, T, D] while PyTorch assumes [T, B, D]
|
|
log_probs = log_probs.transpose(1, 0)
|
|
loss = super().forward(
|
|
log_probs=log_probs, targets=targets, input_lengths=input_lengths, target_lengths=target_lengths
|
|
)
|
|
if self._apply_batch_mean:
|
|
loss = torch.mean(loss)
|
|
return loss
|