From 580b11b57ac4b62f7e2acda73813a4e10e8e4cd3 Mon Sep 17 00:00:00 2001
From: 游雁 <zhifu.gzf@alibaba-inc.com>
Date: 星期二, 10 十月 2023 17:17:29 +0800
Subject: [PATCH] v0.8.0
---
funasr/layers/label_aggregation.py | 39 +++++++++++++++++++++++++++++++++++----
1 files changed, 35 insertions(+), 4 deletions(-)
diff --git a/funasr/layers/label_aggregation.py b/funasr/layers/label_aggregation.py
index 075e19d..e86c296 100644
--- a/funasr/layers/label_aggregation.py
+++ b/funasr/layers/label_aggregation.py
@@ -1,8 +1,7 @@
import torch
-from typeguard import check_argument_types
from typing import Optional
from typing import Tuple
-
+from torch.nn import functional as F
from funasr.modules.nets_utils import make_pad_mask
@@ -13,7 +12,6 @@
hop_length: int = 128,
center: bool = True,
):
- assert check_argument_types()
super().__init__()
self.win_length = win_length
@@ -79,4 +77,37 @@
else:
olens = None
- return output, olens
+ return output.to(input.dtype), olens
+
+
+class LabelAggregateMaxPooling(torch.nn.Module):
+ def __init__(
+ self,
+ hop_length: int = 8,
+ ):
+ super().__init__()
+
+ self.hop_length = hop_length
+
+ def extra_repr(self):
+ return (
+ f"hop_length={self.hop_length}, "
+ )
+
+ def forward(
+ self, input: torch.Tensor, ilens: torch.Tensor = None
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+ """LabelAggregate forward function.
+
+ Args:
+ input: (Batch, Nsamples, Label_dim)
+ ilens: (Batch)
+ Returns:
+ output: (Batch, Frames, Label_dim)
+
+ """
+
+ output = F.max_pool1d(input.transpose(1, 2), self.hop_length, self.hop_length).transpose(1, 2)
+ olens = ilens // self.hop_length
+
+ return output.to(input.dtype), olens
\ No newline at end of file
--
Gitblit v1.9.1