Size the interpolate flops from the whole output tensor

deepspeed/profiling/flops_profiler/profiler.py

@@ -724,24 +724,32 @@ def _upsample_flops_compute(*args, **kwargs):
         size = args[1]
 
     if size is not None:
+        # `size` is the output *spatial* shape only, so the batch and channel dims have to be
+        # carried over from the input. A scalar applies to every spatial dim.
         if isinstance(size, tuple) or isinstance(size, list):
-            return int(_prod(size)), 0
+            output_spatial = _prod(size)
         else:
-            return int(size), 0
+            output_spatial = size**(input.ndim - 2)
+        return int(output_spatial * _prod(input.shape[:2])), 0
 
     scale_factor = kwargs.get('scale_factor', None)
     if scale_factor is None and len(args) > 2:
         scale_factor = args[2]
     assert scale_factor is not None, "either size or scale_factor should be defined"
 
-    flops = input.numel()
     if isinstance(scale_factor, (list, tuple)):
         # see documentation of `F.interpolate`
         # the spatial dims are defined as the last `n-2` dims of the tensor
         assert len(scale_factor) == input.ndim - 2
-        flops *= _prod(scale_factor)
+        scales = scale_factor
     else:
-        flops *= scale_factor**(input.ndim - 2)
+        scales = (scale_factor, ) * (input.ndim - 2)
+
+    # Each output spatial dim is floored on its own, so the scales cannot be multiplied together
+    # and truncated once: `scale_factor=1.4` on a 32x32 input gives 44x44, not floor(32*32*1.96).
+    flops = _prod(input.shape[:2])
+    for dim, scale in zip(input.shape[2:], scales):
+        flops *= int(dim * scale)
 
     return int(flops), 0
 

tests/unit/profiling/flops_profiler/test_flops_profiler.py

@@ -269,6 +269,78 @@ def forward(self, lhs, rhs):
     assert flops == result.numel()
 
 
+@pytest.mark.sequential
+@pytest.mark.parametrize("input_shape, interpolate_kwargs", [
+    ((8, 3, 32, 32), {
+        "size": (64, 64)
+    }),
+    ((8, 3, 32, 32), {
+        "size": 64
+    }),
+    ((8, 3, 32, 32), {
+        "scale_factor": 2
+    }),
+    ((8, 3, 32), {
+        "size": 64
+    }),
+    ((8, 3, 32), {
+        "size": (64, )
+    }),
+    ((2, 3, 8, 8, 8), {
+        "size": 16
+    }),
+    ((2, 3, 8, 8, 8), {
+        "size": (16, 16, 16)
+    }),
+    ((1, 4, 20, 20), {
+        "size": (10, 10)
+    }),
+    ((8, 3, 32, 32), {
+        "scale_factor": 1.4
+    }),
+    ((8, 3, 32, 32), {
+        "scale_factor": (1.4, 2.6)
+    }),
+    ((2, 3, 7, 7, 7), {
+        "scale_factor": 1.5
+    }),
+    ((1, 4, 20, 20), {
+        "scale_factor": (0.7, 1.9)
+    }),
+    ((8, 3, 32, 32), {
+        "scale_factor": 2.0
+    }),
+    ((4, 2, 16), {
+        "scale_factor": 3
+    }),
+    ((1, 1, 5, 5), {
+        "scale_factor": (2, 3)
+    }),
+])
+def test_interpolate_flops(input_shape, interpolate_kwargs):
+    """`size` is the output spatial shape, so the batch and channel dims still multiply it, and a
+    scalar `size` applies to every spatial dim. `scale_factor` sizes the same job from the input:
+    torch floors each output spatial dim on its own, so the scales cannot be multiplied together
+    and truncated once, which only shows up once a scale is not an integer."""
+
+    class Interpolate(torch.nn.Module):
+
+        def forward(self, x):
+            return torch.nn.functional.interpolate(x, **interpolate_kwargs)
+
+    model = Interpolate()
+    x = torch.randn(*input_shape)
+
+    prof = FlopsProfiler(model)
+    prof.start_profile()
+    result = model(x)
+    prof.stop_profile()
+    flops = prof.get_total_flops()
+    prof.end_profile()
+
+    assert flops == result.numel()
+
+
 class Block(torch.nn.Module):
 
     def __init__(self, linear):