From 19467b57f6476cc0ba5493c0dcde3d15a0c88c2c Mon Sep 17 00:00:00 2001
From: zhifu gao <zhifu.gzf@alibaba-inc.com>
Date: 星期一, 27 二月 2023 17:04:19 +0800
Subject: [PATCH] Merge pull request #160 from alibaba-damo-academy/dev_onnx
---
funasr/export/models/modules/decoder_layer.py | 27 +++++++++++++++++++++++++++
1 files changed, 27 insertions(+), 0 deletions(-)
diff --git a/funasr/export/models/modules/decoder_layer.py b/funasr/export/models/modules/decoder_layer.py
index bc306b1..f539452 100644
--- a/funasr/export/models/modules/decoder_layer.py
+++ b/funasr/export/models/modules/decoder_layer.py
@@ -41,3 +41,30 @@
return x, tgt_mask, memory, memory_mask, cache
+
+class DecoderLayer(nn.Module):
+ def __init__(self, model):
+ super().__init__()
+ self.self_attn = model.self_attn
+ self.src_attn = model.src_attn
+ self.feed_forward = model.feed_forward
+ self.norm1 = model.norm1
+ self.norm2 = model.norm2
+ self.norm3 = model.norm3
+
+ def forward(self, tgt, tgt_mask, memory, memory_mask, cache=None):
+ residual = tgt
+ tgt_q = tgt
+ tgt_q_mask = tgt_mask
+ x = residual + self.self_attn(tgt_q, tgt, tgt, tgt_q_mask)
+
+ residual = x
+ x = self.norm2(x)
+
+ x = residual + self.src_attn(x, memory, memory, memory_mask)
+
+ residual = x
+ x = self.norm3(x)
+ x = residual + self.feed_forward(x)
+
+ return x, tgt_mask, memory, memory_mask
--
Gitblit v1.9.1