From 77703ae1fc0a5be2e4b1372df4bd9849d15eefbc Mon Sep 17 00:00:00 2001
From: Zack Ankner <zackankner@gmail.com>
Date: Wed, 10 Feb 2021 21:47:49 -0700
Subject: [PATCH] moving mpp loss into wrapper

---
 vit_pytorch/__init__.py    |  1 -
 vit_pytorch/mpp_loss.py    | 26 --------------------------
 vit_pytorch/mpp_pytorch.py | 24 ++++++++++++++++++++++--
 3 files changed, 22 insertions(+), 29 deletions(-)
 delete mode 100644 vit_pytorch/mpp_loss.py

diff --git a/vit_pytorch/__init__.py b/vit_pytorch/__init__.py
index ad46ad6..1fcf01f 100644
--- a/vit_pytorch/__init__.py
+++ b/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
diff --git a/vit_pytorch/mpp_loss.py b/vit_pytorch/mpp_loss.py
deleted file mode 100644
index e627165..0000000
--- a/vit_pytorch/mpp_loss.py
+++ /dev/null
@@ -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
\ No newline at end of file
diff --git a/vit_pytorch/mpp_pytorch.py b/vit_pytorch/mpp_pytorch.py
index 86b4008..ce37335 100644
--- a/vit_pytorch/mpp_pytorch.py
+++ b/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):