-
Notifications
You must be signed in to change notification settings - Fork 2.2k
/
conformer_modules.py
85 lines (70 loc) · 2.63 KB
/
conformer_modules.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import torch.nn as nn
from nemo.collections.asr.modules.activations import Swish
__all__ = ['ConformerConvolution', 'ConformerFeedForward']
class ConformerConvolution(nn.Module):
"""The convolution module for the Conformer model.
Args:
d_model (int): hidden dimension
kernel_size (int): kernel size for depthwise convolution
"""
def __init__(self, d_model, kernel_size):
super(ConformerConvolution, self).__init__()
assert (kernel_size - 1) % 2 == 0
self.d_model = d_model
self.pointwise_conv1 = nn.Conv1d(
in_channels=d_model, out_channels=d_model * 2, kernel_size=1, stride=1, padding=0, bias=True
)
self.depthwise_conv = nn.Conv1d(
in_channels=d_model,
out_channels=d_model,
kernel_size=kernel_size,
stride=1,
padding=(kernel_size - 1) // 2,
groups=d_model,
bias=True,
)
self.batch_norm = nn.BatchNorm1d(d_model)
self.activation = Swish()
self.pointwise_conv2 = nn.Conv1d(
in_channels=d_model, out_channels=d_model, kernel_size=1, stride=1, padding=0, bias=True
)
def forward(self, x):
x = x.transpose(1, 2)
x = self.pointwise_conv1(x)
x = nn.functional.glu(x, dim=1)
x = self.depthwise_conv(x)
x = self.batch_norm(x)
x = self.activation(x)
x = self.pointwise_conv2(x)
x = x.transpose(1, 2)
return x
class ConformerFeedForward(nn.Module):
"""
feed-forward module of Conformer model.
"""
def __init__(self, d_model, d_ff, dropout, activation=Swish()):
super(ConformerFeedForward, self).__init__()
self.linear1 = nn.Linear(d_model, d_ff)
self.activation = activation
self.dropout = nn.Dropout(p=dropout)
self.linear2 = nn.Linear(d_ff, d_model)
def forward(self, x):
x = self.linear1(x)
x = self.activation(x)
x = self.dropout(x)
x = self.linear2(x)
return x