summaryrefslogtreecommitdiff
path: root/text_recognizer/networks
diff options
context:
space:
mode:
Diffstat (limited to 'text_recognizer/networks')
-rw-r--r--text_recognizer/networks/conformer/conv.py13
-rw-r--r--text_recognizer/networks/conformer/depth_wise_conv.py17
2 files changed, 9 insertions, 21 deletions
diff --git a/text_recognizer/networks/conformer/conv.py b/text_recognizer/networks/conformer/conv.py
index f031dc7..ac13f5d 100644
--- a/text_recognizer/networks/conformer/conv.py
+++ b/text_recognizer/networks/conformer/conv.py
@@ -4,7 +4,6 @@ from einops.layers.torch import Rearrange
from torch import nn, Tensor
-from text_recognizer.networks.conformer.depth_wise_conv import DepthwiseConv1D
from text_recognizer.networks.conformer.glu import GLU
@@ -21,12 +20,18 @@ class ConformerConv(nn.Module):
self.layers = nn.Sequential(
nn.LayerNorm(dim),
Rearrange("b n c -> b c n"),
- nn.Conv1D(dim, 2 * inner_dim, 1),
+ nn.Conv1d(dim, 2 * inner_dim, 1),
GLU(dim=1),
- DepthwiseConv1D(inner_dim, inner_dim, kernel_size),
+ nn.Conv1d(
+ in_channels=inner_dim,
+ out_channels=inner_dim,
+ kernel_size=kernel_size,
+ groups=inner_dim,
+ padding="same",
+ ),
nn.BatchNorm1d(inner_dim),
nn.Mish(inplace=True),
- nn.Conv1D(inner_dim, dim, 1),
+ nn.Conv1d(inner_dim, dim, 1),
Rearrange("b c n -> b n c"),
nn.Dropout(dropout),
)
diff --git a/text_recognizer/networks/conformer/depth_wise_conv.py b/text_recognizer/networks/conformer/depth_wise_conv.py
deleted file mode 100644
index 1dbd0b8..0000000
--- a/text_recognizer/networks/conformer/depth_wise_conv.py
+++ /dev/null
@@ -1,17 +0,0 @@
-"""Depthwise 1D convolution."""
-from torch import nn, Tensor
-
-
-class DepthwiseConv1D(nn.Module):
- def __init__(self, in_channels: int, out_channels: int, kernel_size: int) -> None:
- super().__init__()
- self.conv = nn.Conv1d(
- in_channels=in_channels,
- out_channels=out_channels,
- kernel_size=kernel_size,
- groups=in_channels,
- padding="same",
- )
-
- def forward(self, x: Tensor) -> Tensor:
- return self.conv(x)