#!/usr/bin/env python3
|
# -*- coding: utf-8 -*-
|
|
import torch
|
from torch import nn
|
|
|
class EncoderLayerSANM(nn.Module):
|
def __init__(
|
self,
|
model,
|
):
|
"""Construct an EncoderLayer object."""
|
super().__init__()
|
self.self_attn = model.self_attn
|
self.feed_forward = model.feed_forward
|
self.norm1 = model.norm1
|
self.norm2 = model.norm2
|
self.in_size = model.in_size
|
self.size = model.size
|
|
def forward(self, x, mask):
|
|
residual = x
|
x = self.norm1(x)
|
x = self.self_attn(x, mask)
|
if self.in_size == self.size:
|
x = x + residual
|
residual = x
|
x = self.norm2(x)
|
x = self.feed_forward(x)
|
x = x + residual
|
|
return x, mask
|
|
|
class EncoderLayerConformer(nn.Module):
|
def __init__(
|
self,
|
model,
|
):
|
"""Construct an EncoderLayer object."""
|
super().__init__()
|
self.self_attn = model.self_attn
|
self.feed_forward = model.feed_forward
|
self.feed_forward_macaron = model.feed_forward_macaron
|
self.conv_module = model.conv_module
|
self.norm_ff = model.norm_ff
|
self.norm_mha = model.norm_mha
|
self.norm_ff_macaron = model.norm_ff_macaron
|
self.norm_conv = model.norm_conv
|
self.norm_final = model.norm_final
|
self.size = model.size
|
|
def forward(self, x, mask):
|
if isinstance(x, tuple):
|
x, pos_emb = x[0], x[1]
|
else:
|
x, pos_emb = x, None
|
|
if self.feed_forward_macaron is not None:
|
residual = x
|
x = self.norm_ff_macaron(x)
|
x = residual + self.feed_forward_macaron(x) * 0.5
|
|
residual = x
|
x = self.norm_mha(x)
|
|
x_q = x
|
|
if pos_emb is not None:
|
x_att = self.self_attn(x_q, x, x, pos_emb, mask)
|
else:
|
x_att = self.self_attn(x_q, x, x, mask)
|
x = residual + x_att
|
|
if self.conv_module is not None:
|
residual = x
|
x = self.norm_conv(x)
|
x = residual + self.conv_module(x)
|
|
residual = x
|
x = self.norm_ff(x)
|
x = residual + self.feed_forward(x) * 0.5
|
|
x = self.norm_final(x)
|
|
if pos_emb is not None:
|
return (x, pos_emb), mask
|
|
return x, mask
|
