def input_types(self):
"""Returns definitions of module input ports.
"""
return {
"log_probs": NeuralType(("B", "T", "D"), LogprobsType()),
"log_probs_length": NeuralType(tuple("B"), LengthsType()),
}
def output_types(self):
return {
"outputs": NeuralType(('B', 'T', 'T', 'D'), LogprobsType()),
"prednet_lengths": NeuralType(tuple('B'), LengthsType()),
"output-states-1": NeuralType((('D', 'B', 'D')), ElementType()),
"output-states-2": NeuralType((('D', 'B', 'D')), ElementType()),
}
def output_types(self):
""" Returns definitions of module output ports. """
return {
"predictions":
NeuralType(axes=(AxisType(kind=AxisKind.Batch),
AxisType(kind=AxisKind.Dimension)),
elements_type=LogprobsType())
}
def input_types(self):
""" Returns definitions of module input ports. """
return {
"predictions":
NeuralType(axes=('B', 'ANY'), elements_type=LogprobsType()),
"targets":
NeuralType(axes=('B'), elements_type=ClassificationTarget()),
}
def input_types(self):
"""Returns definitions of module input ports.
"""
return {
"log_probs": NeuralType(("B", "T", "D"), LogprobsType()),
"labels": NeuralType(("B", "T"), LabelsType()),
"output_mask": NeuralType(("B", "T"), MaskType(), optional=True),
}
def output_types(self) -> Optional[Dict[str, NeuralType]]:
"""
Returns definitions of module output ports.
"""
if not self.log_softmax:
return {"logits": NeuralType(('B', 'T', 'C'), LogitsType())}
else:
return {"log_probs": NeuralType(('B', 'T', 'C'), LogprobsType())}
def input_types(self):
"""Input types definitions for LatticeLoss.
"""
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()),
}
def input_types(self):
"""Input types definitions for CTCLoss.
"""
return {
"log_probs": NeuralType(('B', 'T', 'T', 'D'), LogprobsType()),
"targets": NeuralType(('B', 'T'), LabelsType()),
"input_lengths": NeuralType(tuple('B'), LengthsType()),
"target_lengths": NeuralType(tuple('B'), LengthsType()),
}
def output_ports(self):
"""
Returns definitions of module output ports.
"""
# Prepare list of axes.
axes = [AxisType(kind=AxisKind.Batch)]
for size in self._sizes[1:]:
axes.append(AxisType(kind=AxisKind.Any, size=size))
# Return neural type.
# TODO: if self._type != "logsoftmax"
return {"outputs": NeuralType(axes, LogprobsType())}
def input_types(self):
"""Input types definitions for Contrastive.
"""
return {
"spec_masks": NeuralType(("B", "D", "T"), SpectrogramType()),
"decoder_outputs": NeuralType(("B", "T", "D"), LogprobsType()),
"targets": NeuralType(('B', 'T'), LabelsType()),
"decoder_lengths": NeuralType(tuple('B'),
LengthsType(),
optional=True),
"target_lengths": NeuralType(tuple('B'),
LengthsType(),
optional=True),
}
def output_types(self):
"""Returns definitions of module output ports.
"""
if not self._fuse_loss_wer:
return {
"outputs": NeuralType(('B', 'T', 'T', 'D'), LogprobsType()),
}
else:
return {
"loss": NeuralType(elements_type=LossType(), optional=True),
"wer": NeuralType(elements_type=ElementType(), optional=True),
"wer_numer": NeuralType(elements_type=ElementType(), optional=True),
"wer_denom": NeuralType(elements_type=ElementType(), optional=True),
}
def output_types(self) -> Optional[Dict[str, NeuralType]]:
return {
"outputs": NeuralType(('B', 'T', 'D'), LogprobsType()),
"encoded_lengths": NeuralType(tuple('B'), LengthsType()),
"greedy_predictions": NeuralType(('B', 'T'), LabelsType()),
}
def output_types(self):
return OrderedDict(
{"logprobs": NeuralType(('B', 'T', 'D'), LogprobsType())})
def input_ports(self):
"""Returns:
Definitions of module input ports.
"""
return {"log_probs": NeuralType(('B', 'T', 'D'), LogprobsType())}
def output_types(self) -> Optional[Dict[str, NeuralType]]:
if not self.log_softmax:
return {"logits": NeuralType(('B', 'D'), LogitsType())}
else:
return {"log_probs": NeuralType(('B', 'D'), LogprobsType())}
