remove Barlow Twins loss. You another repo if you want to do SSL pretraining. something like LightlySSL for example

Author: pfeatherstone

src/models.py

@@ -1188,27 +1188,4 @@ def conf(): return preds[...,4] if has_objectness else preds[...,4:].max(-1)[0]
     nms     = torchvision.ops.batched_nms(preds[:,:4], conf(), batch, iou_threshold=nms_thresh)
     preds   = preds[nms]
     batch   = batch[nms]
-    return batch, preds
-
-class BarlowTwinsHead(nn.Module):
-    def __init__(self, backbone, input_dim, hidden_dim=2048, output_dim=128):
-        super().__init__()
-        self.net  = backbone
-        self.proj = nn.Sequential(nn.Linear(input_dim, hidden_dim, bias=True),
-                                  nn.LayerNorm(hidden_dim),
-                                  nn.ReLU(),
-                                  nn.Linear(hidden_dim, output_dim, bias=False))
-
-    def forward(self, x):
-        x = self.net(x)[-1]
-        x = x.flatten(2).mean(2)
-        x = self.proj(x)
-        return x
-    
-def barlow_loss(z1, z2, lambda_coeff):
-    z1, z2      = map(lambda z: (z - z.mean(0)) / z.std(0), (z1,z2))
-    cross       = (z1.T @ z2) / z1.shape[0]
-    mask        = torch.eye(cross.shape[0], dtype=torch.bool, device=cross.device)
-    on_diag     = (cross[mask]-1).pow(2).sum()
-    off_diag    = cross[~mask].pow(2).sum()
-    return (on_diag + lambda_coeff * off_diag, cross)
+    return batch, preds