From ca1925433861f6b1037bcd81112d56717d9f153b Mon Sep 17 00:00:00 2001
From: Gustaf Rydholm <gustaf.rydholm@gmail.com>
Date: Mon, 11 Oct 2021 22:13:54 +0200
Subject: Add BarlowNetwork

---
 text_recognizer/networks/barlow_twins/network.py | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)
 create mode 100644 text_recognizer/networks/barlow_twins/network.py

(limited to 'text_recognizer/networks/barlow_twins')

diff --git a/text_recognizer/networks/barlow_twins/network.py b/text_recognizer/networks/barlow_twins/network.py
new file mode 100644
index 0000000..874e570
--- /dev/null
+++ b/text_recognizer/networks/barlow_twins/network.py
@@ -0,0 +1,19 @@
+"""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
+
-- 
cgit v1.2.3-70-g09d2