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

fp16

Browse Source
This commit is contained in:
游雁 2024-06-11 13:49:57 +08:00
parent 23008c7cac
commit 63e60cc43d
2 changed files with 9 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
funasr/auto/auto_model.py
View File

@ -233,6 +233,8 @@ class AutoModel:
# fp16 # fp16
if kwargs.get("fp16", False): if kwargs.get("fp16", False):
model.to(torch.float16) model.to(torch.float16)
elif kwargs.get("bf16", False):
model.to(torch.bfloat16)
return model, kwargs return model, kwargs
def __call__(self, *args, **cfg): def __call__(self, *args, **cfg):

7
funasr/models/llm_asr/model.py
View File

@ -684,6 +684,13 @@ class LLMASR2(nn.Module):
# audio encoder # audio encoder
speech = batch["speech"] speech = batch["speech"]
speech_lengths = batch["speech_lengths"][:, 0] speech_lengths = batch["speech_lengths"][:, 0]
# fp16
if kwargs.get("fp16", False):
speech = speech.to(torch.float16)
encoder_out_lens = encoder_out_lens.to(torch.float16)
elif kwargs.get("bf16", False):
speech = speech.to(torch.bfloat16)
encoder_out_lens = encoder_out_lens.to(torch.bfloat16)
encoder_out, encoder_out_lens = self.audio_encoder(speech.permute(0, 2, 1), speech_lengths) encoder_out, encoder_out_lens = self.audio_encoder(speech.permute(0, 2, 1), speech_lengths)
# audio_adaptor # audio_adaptor