diff --git a/deepspeed/profiling/flops_profiler/profiler.py b/deepspeed/profiling/flops_profiler/profiler.py index a8caaa51fa7e..cec0ca54d95f 100644 --- a/deepspeed/profiling/flops_profiler/profiler.py +++ b/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 diff --git a/tests/unit/profiling/flops_profiler/test_flops_profiler.py b/tests/unit/profiling/flops_profiler/test_flops_profiler.py index 075bfdacd3aa..6c6dbd762ed7 100644 --- a/tests/unit/profiling/flops_profiler/test_flops_profiler.py +++ b/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):