From 4bc6db3ef88795eb570f92f9576f8bc7c56f96bc Mon Sep 17 00:00:00 2001
From: 志浩 <neo.dzh@alibaba-inc.com>
Date: 星期二, 01 八月 2023 17:03:39 +0800
Subject: [PATCH] TOLD: add TOLD/SOND recipe on callhome

---
 funasr/layers/label_aggregation.py |   37 ++++++++++++++++++++++++++++++++++++-
 1 files changed, 36 insertions(+), 1 deletions(-)

diff --git a/funasr/layers/label_aggregation.py b/funasr/layers/label_aggregation.py
index 8366a79..97556d5 100644
--- a/funasr/layers/label_aggregation.py
+++ b/funasr/layers/label_aggregation.py
@@ -1,7 +1,8 @@
 import torch
 from typing import Optional
 from typing import Tuple
-
+from typeguard import check_argument_types
+from torch.nn import functional as F
 from funasr.modules.nets_utils import make_pad_mask
 
 
@@ -78,3 +79,37 @@
             olens = None
 
         return output.to(input.dtype), olens
+
+
+class LabelAggregateMaxPooling(torch.nn.Module):
+    def __init__(
+        self,
+        hop_length: int = 8,
+    ):
+        assert check_argument_types()
+        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