From 1d4ab65c8bfebaecbcb0eec0064bae9a321cad75 Mon Sep 17 00:00:00 2001
From: 游雁 <zhifu.gzf@alibaba-inc.com>
Date: 星期二, 14 二月 2023 16:27:37 +0800
Subject: [PATCH] export model
---
funasr/modules/subsampling.py | 129 +++++++++++++++++++++++++++++++++++++++----
1 files changed, 117 insertions(+), 12 deletions(-)
diff --git a/funasr/modules/subsampling.py b/funasr/modules/subsampling.py
index f9a1c16..d492ccf 100644
--- a/funasr/modules/subsampling.py
+++ b/funasr/modules/subsampling.py
@@ -5,12 +5,12 @@
# Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
"""Subsampling layer definition."""
-
+import numpy as np
import torch
import torch.nn.functional as F
from funasr.modules.embedding import PositionalEncoding
-
-
+import logging
+from funasr.modules.streaming_utils.utils import sequence_mask
class TooShortUttError(Exception):
"""Raised when the utt is too short for subsampling.
@@ -87,6 +87,72 @@
if x_mask is None:
return x, None
return x, x_mask[:, :, :-2:2][:, :, :-2:2]
+
+ def __getitem__(self, key):
+ """Get item.
+
+ When reset_parameters() is called, if use_scaled_pos_enc is used,
+ return the positioning encoding.
+
+ """
+ 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).
+
+ Args:
+ idim (int): Input dimension.
+ odim (int): Output dimension.
+ dropout_rate (float): Dropout rate.
+ pos_enc (torch.nn.Module): Custom position encoding layer.
+
+ """
+
+ def __init__(self, idim, odim, dropout_rate, pos_enc=None):
+ """Construct an Conv2dSubsampling object."""
+ super(Conv2dSubsamplingPad, self).__init__()
+ self.conv = torch.nn.Sequential(
+ torch.nn.Conv2d(1, odim, 3, 2, padding=(0, 0)),
+ torch.nn.ReLU(),
+ torch.nn.Conv2d(odim, odim, 3, 2, padding=(0, 0)),
+ torch.nn.ReLU(),
+ )
+ self.out = torch.nn.Sequential(
+ torch.nn.Linear(odim * (((idim - 1) // 2 - 1) // 2), odim),
+ pos_enc if pos_enc is not None else PositionalEncoding(odim, dropout_rate),
+ )
+ self.pad_fn = torch.nn.ConstantPad1d((0, 4), 0.0)
+
+ def forward(self, x, x_mask):
+ """Subsample x.
+
+ Args:
+ x (torch.Tensor): Input tensor (#batch, time, idim).
+ x_mask (torch.Tensor): Input mask (#batch, 1, time).
+
+ Returns:
+ torch.Tensor: Subsampled tensor (#batch, time', odim),
+ where time' = time // 4.
+ torch.Tensor: Subsampled mask (#batch, 1, time'),
+ where time' = time // 4.
+
+ """
+ x = x.transpose(1, 2)
+ x = self.pad_fn(x)
+ x = x.transpose(1, 2)
+ x = x.unsqueeze(1) # (b, c, t, f)
+ x = self.conv(x)
+ b, c, t, f = x.size()
+ x = self.out(x.transpose(1, 2).contiguous().view(b, t, c * f))
+ if x_mask is None:
+ return x, None
+ x_len = torch.sum(x_mask[:, 0, :], dim=-1)
+ x_len = (x_len - 1) // 2 + 1
+ x_len = (x_len - 1) // 2 + 1
+ mask = sequence_mask(x_len, None, x_len.dtype, x[0].device)
+ return x, mask[:, None, :]
def __getitem__(self, key):
"""Get item.
@@ -267,12 +333,17 @@
"""
- def __init__(self, idim, odim, kernel_size, stride, pad):
+ 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)
self.stride = stride
self.odim = odim
+ self.tf2torch_tensor_name_prefix_torch = tf2torch_tensor_name_prefix_torch
+ self.tf2torch_tensor_name_prefix_tf = tf2torch_tensor_name_prefix_tf
def output_size(self) -> int:
return self.odim
@@ -292,13 +363,47 @@
x_len = (x_len - 1) // self.stride + 1
return x, x_len
- def __getitem__(self, key):
- """Get item.
+ 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
- When reset_parameters() is called, if use_scaled_pos_enc is used,
- return the positioning encoding.
+ 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
- """
- if key != -1:
- raise NotImplementedError("Support only `-1` (for `reset_parameters`).")
- return self.out[key]
--
Gitblit v1.9.1