moving mpp loss into wrapper

vit_pytorch/__init__.py

@@ -1,3 +1,2 @@
 from vit_pytorch.vit_pytorch import ViT
-from vit_pytorch.mpp_loss import MPPLoss
 from vit_pytorch.mpp_pytorch import MPP

vit_pytorch/mpp_loss.py

@@ -1,26 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from einops import rearrange
-
-
-
-class MPPLoss(nn.Module):
-    def __init__(self, patch_size):
-        super(MPPLoss, self).__init__()
-        self.patch_size = patch_size
-
-    def forward(self, predicted_patches, target, mask):
-        # reshape target to patches
-        p = self.patch_size
-        target = rearrange(target, "b c (h p1) (w p2) -> b (h w) c (p1 p2) ", p1 = p, p2 = p)
-
-        channel_bins = torch.tensor([0.333, 0.666, 1.0])
-        target = torch.bucketize(target, channel_bins, right=True)
-        target = target.float().mean(dim=3)
-
-        predicted_patches = predicted_patches[mask]
-        target = target[mask]
-
-        loss = F.mse_loss(predicted_patches, target)
-        return loss

vit_pytorch/mpp_pytorch.py

@@ -7,8 +7,6 @@ import torch.nn.functional as F
 
 from einops import rearrange
 
-from vit_pytorch import MPPLoss
-
 # helpers
 
 def prob_mask_like(t, prob):
@@ -26,6 +24,28 @@ def get_mask_subset_with_prob(patched_input, prob):
     new_mask.scatter_(1, sampled_indices, 1)
     return new_mask.bool()
 
+# mpp loss
+
+class MPPLoss(nn.Module):
+    def __init__(self, patch_size):
+        super(MPPLoss, self).__init__()
+        self.patch_size = patch_size
+
+    def forward(self, predicted_patches, target, mask):
+        # reshape target to patches
+        p = self.patch_size
+        target = rearrange(target, "b c (h p1) (w p2) -> b (h w) c (p1 p2) ", p1 = p, p2 = p)
+
+        channel_bins = torch.tensor([0.333, 0.666, 1.0])
+        target = torch.bucketize(target, channel_bins, right=True)
+        target = target.float().mean(dim=3)
+
+        predicted_patches = predicted_patches[mask]
+        target = target[mask]
+
+        loss = F.mse_loss(predicted_patches, target)
+        return loss
+
 # main class
 
 class MPP(nn.Module):