freedakgmail
6116b1f3c6
Update API Documentation / build-api-docs (push) Has been cancelled
Add complete FunASR codebase including models, runtime, and documentation.
259 lines
8.6 KiB
Python
Executable File
259 lines
8.6 KiB
Python
Executable File
from typing import Tuple, Dict
|
|
import copy
|
|
import os
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from funasr.models.fsmn_kws.encoder import (toKaldiMatrix, LinearTransform, AffineTransform, RectifiedLinear, FSMNBlock, FsmnStack, BasicBlock)
|
|
|
|
|
|
from funasr.register import tables
|
|
|
|
|
|
'''
|
|
FSMN net for keyword spotting
|
|
input_dim: input dimension
|
|
linear_dim: fsmn input dimensionll
|
|
proj_dim: fsmn projection dimension
|
|
lorder: fsmn left order
|
|
rorder: fsmn right order
|
|
num_syn: output dimension
|
|
fsmn_layers: no. of sequential fsmn layers
|
|
'''
|
|
|
|
@tables.register("encoder_classes", "FSMNMT")
|
|
class FSMNMT(nn.Module):
|
|
def __init__(
|
|
self,
|
|
input_dim: int,
|
|
input_affine_dim: int,
|
|
fsmn_layers: int,
|
|
linear_dim: int,
|
|
proj_dim: int,
|
|
lorder: int,
|
|
rorder: int,
|
|
lstride: int,
|
|
rstride: int,
|
|
output_affine_dim: int,
|
|
output_dim: int,
|
|
output_dim2: int,
|
|
use_softmax: bool = True,
|
|
):
|
|
"""Initialize FSMNMT.
|
|
|
|
Args:
|
|
input_dim: Size/dimension parameter.
|
|
input_affine_dim: Size/dimension parameter.
|
|
fsmn_layers: TODO.
|
|
linear_dim: Size/dimension parameter.
|
|
proj_dim: Size/dimension parameter.
|
|
lorder: TODO.
|
|
rorder: TODO.
|
|
lstride: TODO.
|
|
rstride: TODO.
|
|
output_affine_dim: Size/dimension parameter.
|
|
output_dim: Size/dimension parameter.
|
|
output_dim2: TODO.
|
|
use_softmax: TODO.
|
|
"""
|
|
super().__init__()
|
|
|
|
self.input_dim = input_dim
|
|
self.input_affine_dim = input_affine_dim
|
|
self.fsmn_layers = fsmn_layers
|
|
self.linear_dim = linear_dim
|
|
self.proj_dim = proj_dim
|
|
self.output_affine_dim = output_affine_dim
|
|
self.output_dim = output_dim
|
|
self.output_dim2 = output_dim2
|
|
|
|
self.in_linear1 = AffineTransform(input_dim, input_affine_dim)
|
|
self.in_linear2 = AffineTransform(input_affine_dim, linear_dim)
|
|
self.relu = RectifiedLinear(linear_dim, linear_dim)
|
|
self.fsmn = FsmnStack(*[BasicBlock(linear_dim, proj_dim, lorder, rorder, lstride, rstride, i) for i in
|
|
range(fsmn_layers)])
|
|
self.out_linear1 = AffineTransform(linear_dim, output_affine_dim)
|
|
self.out_linear1_2 = AffineTransform(linear_dim, output_affine_dim)
|
|
self.out_linear2 = AffineTransform(output_affine_dim, output_dim)
|
|
self.out_linear2_2 = AffineTransform(output_affine_dim, output_dim2)
|
|
|
|
self.use_softmax = use_softmax
|
|
if self.use_softmax:
|
|
self.softmax = nn.Softmax(dim=-1)
|
|
|
|
def output_size(self) -> int:
|
|
"""Output size."""
|
|
return self.output_dim
|
|
|
|
def output_size2(self) -> int:
|
|
"""Output size2."""
|
|
return self.output_dim2
|
|
|
|
def forward(
|
|
self,
|
|
input: torch.Tensor,
|
|
cache: Dict[str, torch.Tensor] = None
|
|
) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
|
|
"""
|
|
Args:
|
|
input (torch.Tensor): Input tensor (B, T, D)
|
|
cache: when cache is not None, the forward is in streaming. The type of cache is a dict, egs,
|
|
{'cache_layer_1': torch.Tensor(B, T1, D)}, T1 is equal to self.lorder. It is {} for the 1st frame
|
|
"""
|
|
|
|
x1 = self.in_linear1(input)
|
|
x2 = self.in_linear2(x1)
|
|
x3 = self.relu(x2)
|
|
x4 = self.fsmn(x3, cache) # self.cache will update automatically in self.fsmn
|
|
x5 = self.out_linear1(x4)
|
|
x6 = self.out_linear2(x5)
|
|
|
|
x5_2 = self.out_linear1_2(x4)
|
|
x6_2 = self.out_linear2_2(x5_2)
|
|
|
|
if self.use_softmax:
|
|
x7 = self.softmax(x6)
|
|
x7_2 = self.softmax(x6_2)
|
|
return x7, x7_2
|
|
|
|
return x6, x6_2
|
|
|
|
|
|
@tables.register("encoder_classes", "FSMNMTConvert")
|
|
class FSMNMTConvert(nn.Module):
|
|
def __init__(
|
|
self,
|
|
input_dim: int,
|
|
input_affine_dim: int,
|
|
fsmn_layers: int,
|
|
linear_dim: int,
|
|
proj_dim: int,
|
|
lorder: int,
|
|
rorder: int,
|
|
lstride: int,
|
|
rstride: int,
|
|
output_affine_dim: int,
|
|
output_dim: int,
|
|
output_dim2: int,
|
|
use_softmax: bool = True,
|
|
):
|
|
"""Initialize FSMNMTConvert.
|
|
|
|
Args:
|
|
input_dim: Size/dimension parameter.
|
|
input_affine_dim: Size/dimension parameter.
|
|
fsmn_layers: TODO.
|
|
linear_dim: Size/dimension parameter.
|
|
proj_dim: Size/dimension parameter.
|
|
lorder: TODO.
|
|
rorder: TODO.
|
|
lstride: TODO.
|
|
rstride: TODO.
|
|
output_affine_dim: Size/dimension parameter.
|
|
output_dim: Size/dimension parameter.
|
|
output_dim2: TODO.
|
|
use_softmax: TODO.
|
|
"""
|
|
super().__init__()
|
|
|
|
self.input_dim = input_dim
|
|
self.input_affine_dim = input_affine_dim
|
|
self.fsmn_layers = fsmn_layers
|
|
self.linear_dim = linear_dim
|
|
self.proj_dim = proj_dim
|
|
self.output_affine_dim = output_affine_dim
|
|
self.output_dim = output_dim
|
|
self.output_dim2 = output_dim2
|
|
|
|
self.in_linear1 = AffineTransform(input_dim, input_affine_dim)
|
|
self.in_linear2 = AffineTransform(input_affine_dim, linear_dim)
|
|
self.relu = RectifiedLinear(linear_dim, linear_dim)
|
|
self.fsmn = FsmnStack(*[BasicBlock(linear_dim, proj_dim, lorder, rorder, lstride, rstride, i) for i in
|
|
range(fsmn_layers)])
|
|
self.out_linear1 = AffineTransform(linear_dim, output_affine_dim)
|
|
self.out_linear1_2 = AffineTransform(linear_dim, output_affine_dim)
|
|
self.out_linear2 = AffineTransform(output_affine_dim, output_dim)
|
|
self.out_linear2_2 = AffineTransform(output_affine_dim, output_dim2)
|
|
|
|
self.use_softmax = use_softmax
|
|
if self.use_softmax:
|
|
self.softmax = nn.Softmax(dim=-1)
|
|
|
|
def output_size(self) -> int:
|
|
"""Output size."""
|
|
return self.output_dim
|
|
|
|
def output_size2(self) -> int:
|
|
"""Output size2."""
|
|
return self.output_dim2
|
|
|
|
def to_kaldi_net(self):
|
|
"""To kaldi net."""
|
|
re_str = ''
|
|
re_str += '<Nnet>\n'
|
|
re_str += self.in_linear1.to_kaldi_net()
|
|
re_str += self.in_linear2.to_kaldi_net()
|
|
re_str += self.relu.to_kaldi_net()
|
|
|
|
for fsmn in self.fsmn:
|
|
re_str += fsmn.to_kaldi_net()
|
|
|
|
re_str += self.out_linear1.to_kaldi_net()
|
|
re_str += self.out_linear2.to_kaldi_net()
|
|
re_str += '<Softmax> %d %d\n' % (self.output_dim, self.output_dim)
|
|
re_str += '</Nnet>\n'
|
|
|
|
return re_str
|
|
|
|
def to_kaldi_net2(self):
|
|
"""To kaldi net2."""
|
|
re_str = ''
|
|
re_str += '<Nnet>\n'
|
|
re_str += self.in_linear1.to_kaldi_net()
|
|
re_str += self.in_linear2.to_kaldi_net()
|
|
re_str += self.relu.to_kaldi_net()
|
|
|
|
for fsmn in self.fsmn:
|
|
re_str += fsmn.to_kaldi_net()
|
|
|
|
re_str += self.out_linear1_2.to_kaldi_net()
|
|
re_str += self.out_linear2_2.to_kaldi_net()
|
|
re_str += '<Softmax> %d %d\n' % (self.output_dim2, self.output_dim2)
|
|
re_str += '</Nnet>\n'
|
|
|
|
return re_str
|
|
|
|
def to_pytorch_net(self, kaldi_file):
|
|
"""To pytorch net.
|
|
|
|
Args:
|
|
kaldi_file: TODO.
|
|
"""
|
|
with open(kaldi_file, 'r', encoding='utf8') as fread:
|
|
fread = open(kaldi_file, 'r')
|
|
nnet_start_line = fread.readline()
|
|
assert nnet_start_line.strip() == '<Nnet>'
|
|
|
|
self.in_linear1.to_pytorch_net(fread)
|
|
self.in_linear2.to_pytorch_net(fread)
|
|
self.relu.to_pytorch_net(fread)
|
|
|
|
for fsmn in self.fsmn:
|
|
fsmn.to_pytorch_net(fread)
|
|
|
|
self.out_linear1.to_pytorch_net(fread)
|
|
self.out_linear2.to_pytorch_net(fread)
|
|
|
|
softmax_line = fread.readline()
|
|
softmax_split = softmax_line.strip().split()
|
|
assert softmax_split[0].strip() == '<Softmax>'
|
|
assert int(softmax_split[1]) == self.output_dim
|
|
assert int(softmax_split[2]) == self.output_dim
|
|
|
|
nnet_end_line = fread.readline()
|
|
assert nnet_end_line.strip() == '</Nnet>'
|
|
fread.close()
|