mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-18 07:41:03 +00:00
Brandon Thomas
96 lines
3.3 KiB
Python
96 lines
3.3 KiB
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.
|
|
|
|
from collections import OrderedDict
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from nemo.core.classes.common import typecheck
|
|
from nemo.core.classes.exportable import Exportable
|
|
from nemo.core.classes.module import NeuralModule
|
|
from nemo.core.neural_types import AcousticEncodedRepresentation, LogprobsType, NeuralType
|
|
|
|
__all__ = ['LSTMDecoder']
|
|
|
|
|
|
class LSTMDecoder(NeuralModule, Exportable):
|
|
"""
|
|
Simple LSTM Decoder for ASR models
|
|
Args:
|
|
feat_in (int): size of the input features
|
|
num_classes (int): the size of the vocabulary
|
|
lstm_hidden_size (int): hidden size of the LSTM layers
|
|
vocabulary (vocab): The vocabulary
|
|
bidirectional (bool): default is False. Whether LSTMs are bidirectional or not
|
|
num_layers (int): default is 1. Number of LSTM layers stacked
|
|
"""
|
|
|
|
@property
|
|
def input_types(self):
|
|
return OrderedDict({"encoder_output": NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation())})
|
|
|
|
@property
|
|
def output_types(self):
|
|
return OrderedDict({"logprobs": NeuralType(('B', 'T', 'D'), LogprobsType())})
|
|
|
|
def __init__(self, feat_in, num_classes, lstm_hidden_size, vocabulary=None, bidirectional=False, num_layers=1):
|
|
super().__init__()
|
|
|
|
if vocabulary is not None:
|
|
if num_classes != len(vocabulary):
|
|
raise ValueError(
|
|
f"If vocabulary is specified, it's length should be equal to the num_classes. "
|
|
f"Instead got: num_classes={num_classes} and len(vocabulary)={len(vocabulary)}"
|
|
)
|
|
self.__vocabulary = vocabulary
|
|
self._feat_in = feat_in
|
|
# Add 1 for blank char
|
|
self._num_classes = num_classes + 1
|
|
|
|
self.lstm_layer = nn.LSTM(
|
|
input_size=feat_in,
|
|
hidden_size=lstm_hidden_size,
|
|
num_layers=num_layers,
|
|
batch_first=True,
|
|
bidirectional=bidirectional,
|
|
)
|
|
self.linear_layer = torch.nn.Linear(in_features=lstm_hidden_size, out_features=self._num_classes)
|
|
|
|
@typecheck()
|
|
def forward(self, encoder_output):
|
|
output = encoder_output.transpose(1, 2)
|
|
output, _ = self.lstm_layer(output)
|
|
output = self.linear_layer(output)
|
|
return torch.nn.functional.log_softmax(output, dim=-1)
|
|
|
|
def input_example(self):
|
|
"""
|
|
Generates input examples for tracing etc.
|
|
Returns:
|
|
A tuple of input examples.
|
|
"""
|
|
bs = 8
|
|
seq = 64
|
|
input_example = torch.randn(bs, self._feat_in, seq).to(next(self.parameters()).device)
|
|
return tuple([input_example])
|
|
|
|
@property
|
|
def vocabulary(self):
|
|
return self.__vocabulary
|
|
|
|
@property
|
|
def num_classes_with_blank(self):
|
|
return self._num_classes
|