From 38dc6ca3b787bcdb54d43ac5c076e08af25d44b2 Mon Sep 17 00:00:00 2001 From: Gustaf Rydholm Date: Tue, 7 Jun 2022 00:24:28 +0200 Subject: Add subsampler layer --- text_recognizer/networks/conformer/conformer.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) (limited to 'text_recognizer/networks/conformer/conformer.py') diff --git a/text_recognizer/networks/conformer/conformer.py b/text_recognizer/networks/conformer/conformer.py index d56955e..8d0e98e 100644 --- a/text_recognizer/networks/conformer/conformer.py +++ b/text_recognizer/networks/conformer/conformer.py @@ -1,5 +1,6 @@ """Conformer module.""" from copy import deepcopy +from typing import Type from torch import nn, Tensor @@ -7,11 +8,18 @@ from text_recognizer.networks.conformer.block import ConformerBlock class Conformer(nn.Module): - def __init__(self, block: ConformerBlock, depth: int) -> None: + def __init__( + self, + subsampler: Type[nn.Module], + block: ConformerBlock, + depth: int, + ) -> None: super().__init__() + self.subsampler = subsampler self.blocks = nn.ModuleList([deepcopy(block) for _ in range(depth)]) def forward(self, x: Tensor) -> Tensor: + x = self.subsampler(x) for fn in self.blocks: x = fn(x) return x -- cgit v1.2.3-70-g09d2