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 | 35 ++++++++++++++++++++++++++++++++++-
1 files changed, 34 insertions(+), 1 deletions(-)
diff --git a/funasr/layers/label_aggregation.py b/funasr/layers/label_aggregation.py
index 8366a79..e86c296 100644
--- a/funasr/layers/label_aggregation.py
+++ b/funasr/layers/label_aggregation.py
@@ -1,7 +1,7 @@
import torch
from typing import Optional
from typing import Tuple
-
+from torch.nn import functional as F
from funasr.modules.nets_utils import make_pad_mask
@@ -78,3 +78,36 @@
olens = None
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