train on EPE sum instead of mean EPE

Clément Pinard · Clément Pinard · commit f450bdb5976a · 2017-11-29T12:06:55.000+01:00
diff --git a/multiscaleloss.py b/multiscaleloss.py
@@ -2,12 +2,14 @@
 import torch.nn as nn
 
 
-def EPE(input_flow, target_flow, sparse=False):
+def EPE(input_flow, target_flow, sparse=False, mean=True):
     EPE_map = torch.norm(target_flow-input_flow,2,1)
     if sparse:
-        return EPE_map[target_flow == 0].mean()
-    else:
+        EPE_map = EPE_map[target_flow == 0]
+    if mean:
         return EPE_map.mean()
+    else:
+        return EPE_map.sum()
 
 
 def multiscaleEPE(network_output, target_flow, weights=None, sparse=False):
@@ -19,7 +21,7 @@ def one_scale(output, target, sparse):
             target_scaled = nn.functional.adaptive_max_pool2d(target, (h, w))
         else:
             target_scaled = nn.functional.adaptive_avg_pool2d(target, (h, w))
-        return EPE(output, target_scaled, sparse)
+        return EPE(output, target_scaled, sparse, mean=False)
 
     if type(network_output) not in [tuple, list]:
         network_output = [network_output]
@@ -36,4 +38,4 @@ def one_scale(output, target, sparse):
 def realEPE(output, target, sparse=False):
     b, _, h, w = target.size()
     upsampled_output = nn.functional.upsample(output, size=(h,w), mode='bilinear')
-    return EPE(upsampled_output, target, sparse)
+    return EPE(upsampled_output, target, sparse, mean=True)
