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