speech_recognition/FunASR
1
0
Fork 0
mirror of https://github.com/modelscope/FunASR synced 2025-09-15 14:48:36 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

onnx

Browse Source
This commit is contained in:
游雁 2023-04-07 14:37:30 +08:00
parent 9f6445d39b
commit b9837bfc73
5 changed files with 95 additions and 112 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
funasr/export/models/__init__.py
View File

@ -4,10 +4,10 @@ from funasr.export.models.e2e_asr_paraformer import BiCifParaformer as BiCifPara
from funasr.models.e2e_vad import E2EVadModel
from funasr.export.models.e2e_vad import E2EVadModel as E2EVadModel_export
from funasr.models.target_delay_transformer import TargetDelayTransformer
from funasr.export.models.target_delay_transformer import TargetDelayTransformer as TargetDelayTransformer_export
from funasr.export.models.target_delay_transformer import CT_Transformer as CT_Transformer_export
from funasr.train.abs_model import PunctuationModel
from funasr.models.vad_realtime_transformer import VadRealtimeTransformer
from funasr.export.models.vad_realtime_transformer import VadRealtimeTransformer as VadRealtimeTransformer_export
from funasr.export.models.target_delay_transformer import CT_Transformer_VadRealtime as CT_Transformer_VadRealtime_export
def get_model(model, export_config=None):
if isinstance(model, BiCifParaformer):
@ -18,8 +18,8 @@ def get_model(model, export_config=None):
return E2EVadModel_export(model, **export_config)
elif isinstance(model, PunctuationModel):
if isinstance(model.punc_model, TargetDelayTransformer):
return TargetDelayTransformer_export(model.punc_model, **export_config)
return CT_Transformer_export(model.punc_model, **export_config)
elif isinstance(model.punc_model, VadRealtimeTransformer):
return VadRealtimeTransformer_export(model.punc_model, **export_config)
return CT_Transformer_VadRealtime_export(model.punc_model, **export_config)
else:
raise "Funasr does not support the given model type currently."

97
funasr/export/models/target_delay_transformer.py
View File

@ -3,7 +3,12 @@ from typing import Tuple
import torch
import torch.nn as nn
class TargetDelayTransformer(nn.Module):
from funasr.models.encoder.sanm_encoder import SANMEncoder
from funasr.export.models.encoder.sanm_encoder import SANMEncoder as SANMEncoder_export
from funasr.models.encoder.sanm_encoder import SANMVadEncoder
from funasr.export.models.encoder.sanm_encoder import SANMVadEncoder as SANMVadEncoder_export
class CT_Transformer(nn.Module):
def __init__(
self,
@ -23,16 +28,12 @@ class TargetDelayTransformer(nn.Module):
self.num_embeddings = self.embed.num_embeddings
self.model_name = model_name
# from funasr.models.encoder.sanm_encoder import SANMEncoder as Encoder
from funasr.models.encoder.sanm_encoder import SANMEncoder
from funasr.export.models.encoder.sanm_encoder import SANMEncoder as SANMEncoder_export
if isinstance(model.encoder, SANMEncoder):
self.encoder = SANMEncoder_export(model.encoder, onnx=onnx)
else:
assert False, "Only support samn encode."
def forward(self, input: torch.Tensor, text_lengths: torch.Tensor) -> Tuple[torch.Tensor, None]:
def forward(self, inputs: torch.Tensor, text_lengths: torch.Tensor) -> Tuple[torch.Tensor, None]:
"""Compute loss value from buffer sequences.
Args:
@ -40,7 +41,7 @@ class TargetDelayTransformer(nn.Module):
hidden (torch.Tensor): Target ids. (batch, len)
"""
x = self.embed(input)
x = self.embed(inputs)
# mask = self._target_mask(input)
h, _ = self.encoder(x, text_lengths)
y = self.decoder(h)
@ -53,14 +54,14 @@ class TargetDelayTransformer(nn.Module):
return (text_indexes, text_lengths)
def get_input_names(self):
return ['input', 'text_lengths']
return ['inputs', 'text_lengths']
def get_output_names(self):
return ['logits']
def get_dynamic_axes(self):
return {
'input': {
'inputs': {
0: 'batch_size',
1: 'feats_length'
},
@ -73,3 +74,81 @@ class TargetDelayTransformer(nn.Module):
},
}
class CT_Transformer_VadRealtime(nn.Module):
def __init__(
self,
model,
max_seq_len=512,
model_name='punc_model',
**kwargs,
):
super().__init__()
onnx = False
if "onnx" in kwargs:
onnx = kwargs["onnx"]
self.embed = model.embed
if isinstance(model.encoder, SANMVadEncoder):
self.encoder = SANMVadEncoder_export(model.encoder, onnx=onnx)
else:
assert False, "Only support samn encode."
self.decoder = model.decoder
self.model_name = model_name
def forward(self, inputs: torch.Tensor,
text_lengths: torch.Tensor,
vad_indexes: torch.Tensor,
sub_masks: torch.Tensor,
) -> Tuple[torch.Tensor, None]:
"""Compute loss value from buffer sequences.
Args:
input (torch.Tensor): Input ids. (batch, len)
hidden (torch.Tensor): Target ids. (batch, len)
"""
x = self.embed(inputs)
# mask = self._target_mask(input)
h, _ = self.encoder(x, text_lengths, vad_indexes, sub_masks)
y = self.decoder(h)
return y
def with_vad(self):
return True
def get_dummy_inputs(self):
length = 120
text_indexes = torch.randint(0, self.embed.num_embeddings, (1, length))
text_lengths = torch.tensor([length], dtype=torch.int32)
vad_mask = torch.ones(length, length, dtype=torch.float32)[None, None, :, :]
sub_masks = torch.ones(length, length, dtype=torch.float32)
sub_masks = torch.tril(sub_masks).type(torch.float32)
return (text_indexes, text_lengths, vad_mask, sub_masks[None, None, :, :])
def get_input_names(self):
return ['inputs', 'text_lengths', 'vad_masks', 'sub_masks']
def get_output_names(self):
return ['logits']
def get_dynamic_axes(self):
return {
'inputs': {
1: 'feats_length'
},
'vad_masks': {
2: 'feats_length1',
3: 'feats_length2'
},
'sub_masks': {
2: 'feats_length1',
3: 'feats_length2'
},
'logits': {
1: 'logits_length'
},
}

96
funasr/export/models/vad_realtime_transformer.py
View File

@ -1,96 +0,0 @@
from typing import Tuple
import torch
import torch.nn as nn
from funasr.models.encoder.sanm_encoder import SANMVadEncoder
from funasr.export.models.encoder.sanm_encoder import SANMVadEncoder as SANMVadEncoder_export
class VadRealtimeTransformer(nn.Module):
def __init__(
self,
model,
max_seq_len=512,
model_name='punc_model',
**kwargs,
):
super().__init__()
onnx = False
if "onnx" in kwargs:
onnx = kwargs["onnx"]
self.embed = model.embed
if isinstance(model.encoder, SANMVadEncoder):
self.encoder = SANMVadEncoder_export(model.encoder, onnx=onnx)
else:
assert False, "Only support samn encode."
# self.encoder = model.encoder
self.decoder = model.decoder
self.model_name = model_name
def forward(self, input: torch.Tensor,
text_lengths: torch.Tensor,
vad_indexes: torch.Tensor,
sub_masks: torch.Tensor,
) -> Tuple[torch.Tensor, None]:
"""Compute loss value from buffer sequences.
Args:
input (torch.Tensor): Input ids. (batch, len)
hidden (torch.Tensor): Target ids. (batch, len)
"""
x = self.embed(input)
# mask = self._target_mask(input)
h, _ = self.encoder(x, text_lengths, vad_indexes, sub_masks)
y = self.decoder(h)
return y
def with_vad(self):
return True
# def get_dummy_inputs(self):
# length = 120
# text_indexes = torch.randint(0, self.embed.num_embeddings, (1, length))
# text_lengths = torch.tensor([length], dtype=torch.int32)
# vad_mask = torch.ones(length, length, dtype=torch.float32)[None, None, :, :]
# sub_masks = torch.ones(length, length, dtype=torch.float32)
# sub_masks = torch.tril(sub_masks).type(torch.float32)
# return (text_indexes, text_lengths, vad_mask, sub_masks[None, None, :, :])
def get_dummy_inputs(self, txt_dir=None):
from funasr.modules.mask import vad_mask
length = 10
text_indexes = torch.tensor([[266757, 266757, 266757, 266757, 266757, 266757, 266757, 266757, 266757, 266757]], dtype=torch.int32)
text_lengths = torch.tensor([length], dtype=torch.int32)
vad_masks = vad_mask(10, 2, dtype=torch.float32)[None, None, :, :]
sub_masks = torch.ones(length, length, dtype=torch.float32)
sub_masks = torch.tril(sub_masks).type(torch.float32)
return (text_indexes, text_lengths, vad_masks, sub_masks[None, None, :, :])
def get_input_names(self):
return ['input', 'text_lengths', 'vad_masks', 'sub_masks']
def get_output_names(self):
return ['logits']
def get_dynamic_axes(self):
return {
'input': {
1: 'feats_length'
},
'vad_masks': {
2: 'feats_length1',
3: 'feats_length2'
},
'sub_masks': {
2: 'feats_length1',
3: 'feats_length2'
},
'logits': {
1: 'logits_length'
},
}

2
funasr/export/test/test_onnx_punc.py
View File

@ -9,7 +9,7 @@ if __name__ == '__main__':
output_name = [nd.name for nd in sess.get_outputs()]
def _get_feed_dict(text_length):
return {'input': np.ones((1, text_length), dtype=np.int64), 'text_lengths': np.array([text_length,], dtype=np.int32)}
return {'inputs': np.ones((1, text_length), dtype=np.int64), 'text_lengths': np.array([text_length,], dtype=np.int32)}
def _run(feed_dict):
output = sess.run(output_name, input_feed=feed_dict)

4
funasr/export/test/test_onnx_punc_vadrealtime.py
View File

@ -9,9 +9,9 @@ if __name__ == '__main__':
output_name = [nd.name for nd in sess.get_outputs()]
def _get_feed_dict(text_length):
return {'input': np.ones((1, text_length), dtype=np.int64),
return {'inputs': np.ones((1, text_length), dtype=np.int64),
'text_lengths': np.array([text_length,], dtype=np.int32),
'vad_mask': np.ones((1, 1, text_length, text_length), dtype=np.float32),
'vad_masks': np.ones((1, 1, text_length, text_length), dtype=np.float32),
'sub_masks': np.tril(np.ones((text_length, text_length), dtype=np.float32))[None, None, :, :].astype(np.float32)
}