From 28ccfbfc51068a663a80764e14074df5edf2b5ba Mon Sep 17 00:00:00 2001
From: kongdeqiang <kongdeqiang960204@163.com>
Date: 星期五, 13 三月 2026 17:41:41 +0800
Subject: [PATCH] 提交
---
funasr/models/transformer/utils/subsampling.py | 160 +++++++++++++++++++++++++++++------------------------
1 files changed, 88 insertions(+), 72 deletions(-)
diff --git a/funasr/models/transformer/utils/subsampling.py b/funasr/models/transformer/utils/subsampling.py
index 64a9dbc..b18bd31 100644
--- a/funasr/models/transformer/utils/subsampling.py
+++ b/funasr/models/transformer/utils/subsampling.py
@@ -15,6 +15,7 @@
from typing import Optional, Tuple, Union
import math
+
class TooShortUttError(Exception):
"""Raised when the utt is too short for subsampling.
@@ -102,6 +103,7 @@
if key != -1:
raise NotImplementedError("Support only `-1` (for `reset_parameters`).")
return self.out[key]
+
class Conv2dSubsamplingPad(torch.nn.Module):
"""Convolutional 2D subsampling (to 1/4 length).
@@ -326,6 +328,7 @@
return x, None
return x, x_mask[:, :, :-2:2][:, :, :-2:2][:, :, :-2:2]
+
class Conv1dSubsampling(torch.nn.Module):
"""Convolutional 1D subsampling (to 1/2 length).
@@ -337,10 +340,16 @@
"""
- def __init__(self, idim, odim, kernel_size, stride, pad,
- tf2torch_tensor_name_prefix_torch: str = "stride_conv",
- tf2torch_tensor_name_prefix_tf: str = "seq2seq/proj_encoder/downsampling",
- ):
+ def __init__(
+ self,
+ idim,
+ odim,
+ kernel_size,
+ stride,
+ pad,
+ tf2torch_tensor_name_prefix_torch: str = "stride_conv",
+ tf2torch_tensor_name_prefix_tf: str = "seq2seq/proj_encoder/downsampling",
+ ):
super(Conv1dSubsampling, self).__init__()
self.conv = torch.nn.Conv1d(idim, odim, kernel_size, stride)
self.pad_fn = torch.nn.ConstantPad1d(pad, 0.0)
@@ -353,13 +362,11 @@
return self.odim
def forward(self, x, x_len):
- """Subsample x.
-
- """
+ """Subsample x."""
x = x.transpose(1, 2) # (b, d ,t)
x = self.pad_fn(x)
- #x = F.relu(self.conv(x))
- x = F.leaky_relu(self.conv(x), negative_slope=0.)
+ # x = F.relu(self.conv(x))
+ x = F.leaky_relu(self.conv(x), negative_slope=0.0)
x = x.transpose(1, 2) # (b, t ,d)
if x_len is None:
@@ -368,49 +375,6 @@
x_len = (x_len - 1) // self.stride + 1
return x, x_len
- def gen_tf2torch_map_dict(self):
- tensor_name_prefix_torch = self.tf2torch_tensor_name_prefix_torch
- tensor_name_prefix_tf = self.tf2torch_tensor_name_prefix_tf
- map_dict_local = {
- ## predictor
- "{}.conv.weight".format(tensor_name_prefix_torch):
- {"name": "{}/conv1d/kernel".format(tensor_name_prefix_tf),
- "squeeze": None,
- "transpose": (2, 1, 0),
- }, # (256,256,3),(3,256,256)
- "{}.conv.bias".format(tensor_name_prefix_torch):
- {"name": "{}/conv1d/bias".format(tensor_name_prefix_tf),
- "squeeze": None,
- "transpose": None,
- }, # (256,),(256,)
- }
- return map_dict_local
-
- def convert_tf2torch(self,
- var_dict_tf,
- var_dict_torch,
- ):
-
- map_dict = self.gen_tf2torch_map_dict()
-
- var_dict_torch_update = dict()
- for name in sorted(var_dict_torch.keys(), reverse=False):
- names = name.split('.')
- if names[0] == self.tf2torch_tensor_name_prefix_torch:
- name_tf = map_dict[name]["name"]
- data_tf = var_dict_tf[name_tf]
- if map_dict[name]["squeeze"] is not None:
- data_tf = np.squeeze(data_tf, axis=map_dict[name]["squeeze"])
- if map_dict[name]["transpose"] is not None:
- data_tf = np.transpose(data_tf, map_dict[name]["transpose"])
- data_tf = torch.from_numpy(data_tf).type(torch.float32).to("cpu")
-
- var_dict_torch_update[name] = data_tf
-
- logging.info(
- "torch tensor: {}, {}, loading from tf tensor: {}, {}".format(name, data_tf.size(), name_tf,
- var_dict_tf[name_tf].shape))
- return var_dict_torch_update
class StreamingConvInput(torch.nn.Module):
"""Streaming ConvInput module definition.
@@ -438,14 +402,38 @@
conv_size1, conv_size2 = conv_size
self.conv = torch.nn.Sequential(
- torch.nn.Conv2d(1, conv_size1, conv_kernel_size, stride=1, padding=(conv_kernel_size-1)//2),
+ torch.nn.Conv2d(
+ 1,
+ conv_size1,
+ conv_kernel_size,
+ stride=1,
+ padding=(conv_kernel_size - 1) // 2,
+ ),
torch.nn.ReLU(),
- torch.nn.Conv2d(conv_size1, conv_size1, conv_kernel_size, stride=1, padding=(conv_kernel_size-1)//2),
+ torch.nn.Conv2d(
+ conv_size1,
+ conv_size1,
+ conv_kernel_size,
+ stride=1,
+ padding=(conv_kernel_size - 1) // 2,
+ ),
torch.nn.ReLU(),
torch.nn.MaxPool2d((1, 2)),
- torch.nn.Conv2d(conv_size1, conv_size2, conv_kernel_size, stride=1, padding=(conv_kernel_size-1)//2),
+ torch.nn.Conv2d(
+ conv_size1,
+ conv_size2,
+ conv_kernel_size,
+ stride=1,
+ padding=(conv_kernel_size - 1) // 2,
+ ),
torch.nn.ReLU(),
- torch.nn.Conv2d(conv_size2, conv_size2, conv_kernel_size, stride=1, padding=(conv_kernel_size-1)//2),
+ torch.nn.Conv2d(
+ conv_size2,
+ conv_size2,
+ conv_kernel_size,
+ stride=1,
+ padding=(conv_kernel_size - 1) // 2,
+ ),
torch.nn.ReLU(),
torch.nn.MaxPool2d((1, 2)),
)
@@ -464,14 +452,38 @@
kernel_1 = int(subsampling_factor / 2)
self.conv = torch.nn.Sequential(
- torch.nn.Conv2d(1, conv_size1, conv_kernel_size, stride=1, padding=(conv_kernel_size-1)//2),
+ torch.nn.Conv2d(
+ 1,
+ conv_size1,
+ conv_kernel_size,
+ stride=1,
+ padding=(conv_kernel_size - 1) // 2,
+ ),
torch.nn.ReLU(),
- torch.nn.Conv2d(conv_size1, conv_size1, conv_kernel_size, stride=1, padding=(conv_kernel_size-1)//2),
+ torch.nn.Conv2d(
+ conv_size1,
+ conv_size1,
+ conv_kernel_size,
+ stride=1,
+ padding=(conv_kernel_size - 1) // 2,
+ ),
torch.nn.ReLU(),
torch.nn.MaxPool2d((kernel_1, 2)),
- torch.nn.Conv2d(conv_size1, conv_size2, conv_kernel_size, stride=1, padding=(conv_kernel_size-1)//2),
+ torch.nn.Conv2d(
+ conv_size1,
+ conv_size2,
+ conv_kernel_size,
+ stride=1,
+ padding=(conv_kernel_size - 1) // 2,
+ ),
torch.nn.ReLU(),
- torch.nn.Conv2d(conv_size2, conv_size2, conv_kernel_size, stride=1, padding=(conv_kernel_size-1)//2),
+ torch.nn.Conv2d(
+ conv_size2,
+ conv_size2,
+ conv_kernel_size,
+ stride=1,
+ padding=(conv_kernel_size - 1) // 2,
+ ),
torch.nn.ReLU(),
torch.nn.MaxPool2d((2, 2)),
)
@@ -487,9 +499,9 @@
else:
if subsampling_factor == 1:
self.conv = torch.nn.Sequential(
- torch.nn.Conv2d(1, conv_size, 3, [1,2], [1,0]),
+ torch.nn.Conv2d(1, conv_size, 3, [1, 2], [1, 0]),
torch.nn.ReLU(),
- torch.nn.Conv2d(conv_size, conv_size, conv_kernel_size, [1,2], [1,0]),
+ torch.nn.Conv2d(conv_size, conv_size, conv_kernel_size, [1, 2], [1, 0]),
torch.nn.ReLU(),
)
@@ -508,9 +520,11 @@
)
self.conv = torch.nn.Sequential(
- torch.nn.Conv2d(1, conv_size, 3, 2, [1,0]),
+ torch.nn.Conv2d(1, conv_size, 3, 2, [1, 0]),
torch.nn.ReLU(),
- torch.nn.Conv2d(conv_size, conv_size, kernel_2, stride_2, [(kernel_2-1)//2, 0]),
+ torch.nn.Conv2d(
+ conv_size, conv_size, kernel_2, stride_2, [(kernel_2 - 1) // 2, 0]
+ ),
torch.nn.ReLU(),
)
@@ -548,30 +562,32 @@
olens = max(mask.eq(0).sum(1))
b, t, f = x.size()
- x = x.unsqueeze(1) # (b. 1. t. f)
+ x = x.unsqueeze(1) # (b. 1. t. f)
if chunk_size is not None:
max_input_length = int(
- chunk_size * self.subsampling_factor * (math.ceil(float(t) / (chunk_size * self.subsampling_factor) ))
+ chunk_size
+ * self.subsampling_factor
+ * (math.ceil(float(t) / (chunk_size * self.subsampling_factor)))
)
x = map(lambda inputs: pad_to_len(inputs, max_input_length, 1), x)
x = list(x)
x = torch.stack(x, dim=0)
- N_chunks = max_input_length // ( chunk_size * self.subsampling_factor)
+ N_chunks = max_input_length // (chunk_size * self.subsampling_factor)
x = x.view(b * N_chunks, 1, chunk_size * self.subsampling_factor, f)
x = self.conv(x)
_, c, _, f = x.size()
if chunk_size is not None:
- x = x.transpose(1, 2).contiguous().view(b, -1, c * f)[:,:olens,:]
+ x = x.transpose(1, 2).contiguous().view(b, -1, c * f)[:, :olens, :]
else:
x = x.transpose(1, 2).contiguous().view(b, -1, c * f)
if self.output is not None:
x = self.output(x)
- return x, mask[:,:olens][:,:x.size(1)]
+ return x, mask[:, :olens][:, : x.size(1)]
def create_new_vgg_mask(self, mask: torch.Tensor) -> torch.Tensor:
"""Create a new mask for VGG output sequences.
@@ -581,8 +597,8 @@
mask: Mask of output sequences. (B, sub(T))
"""
if self.subsampling_factor > 1:
- vgg1_t_len = mask.size(1) - (mask.size(1) % (self.subsampling_factor // 2 ))
- mask = mask[:, :vgg1_t_len][:, ::self.subsampling_factor // 2]
+ vgg1_t_len = mask.size(1) - (mask.size(1) % (self.subsampling_factor // 2))
+ mask = mask[:, :vgg1_t_len][:, :: self.subsampling_factor // 2]
vgg2_t_len = mask.size(1) - (mask.size(1) % 2)
mask = mask[:, :vgg2_t_len][:, ::2]
@@ -599,7 +615,7 @@
mask: Mask of output sequences. (B, sub(T))
"""
if self.subsampling_factor > 1:
- return mask[:, ::2][:, ::self.stride_2]
+ return mask[:, ::2][:, :: self.stride_2]
else:
return mask
--
Gitblit v1.9.1