diff options
Diffstat (limited to 'text_recognizer/networks/barlow_twins/network.py')
-rw-r--r-- | text_recognizer/networks/barlow_twins/network.py | 18 |
1 files changed, 0 insertions, 18 deletions
diff --git a/text_recognizer/networks/barlow_twins/network.py b/text_recognizer/networks/barlow_twins/network.py deleted file mode 100644 index a3e3750..0000000 --- a/text_recognizer/networks/barlow_twins/network.py +++ /dev/null @@ -1,18 +0,0 @@ -"""Barlow Twins network.""" -from typing import Type - -from torch import nn, Tensor -import torch.nn.functional as F - - -class BarlowTwins(nn.Module): - def __init__(self, encoder: Type[nn.Module], projector: Type[nn.Module]) -> None: - super().__init__() - self.encoder = encoder - self.projector = projector - - def forward(self, x: Tensor) -> Tensor: - z = self.encoder(x) - z_e = F.adaptive_avg_pool2d(z, (1, 1)).flatten(start_dim=1) - z_p = self.projector(z_e) - return z_p |