Fix flops profiler counts for transposed convolutions

deepspeed/profiling/flops_profiler/profiler.py

@@ -623,7 +623,12 @@ def _conv_trans_flops_compute(
 ):
     batch_size = input.shape[0]
     in_channels = input.shape[1]
-    out_channels = weight.shape[1]
+    # A transposed-convolution weight is [in_channels, out_channels // groups, *kernel],
+    # so dim 1 already holds the per-group filter count and must not be divided by groups
+    # a second time. Doing so reports zero macs for a depthwise transposed convolution,
+    # where out_channels // groups is 1.
+    filters_per_channel = weight.shape[1]
+    out_channels = filters_per_channel * groups
     kernel_dims = list(weight.shape[2:])
     input_dims = list(input.shape[2:])
 
@@ -632,19 +637,19 @@ def _conv_trans_flops_compute(
     paddings = padding if type(padding) is tuple else (padding, ) * length
     strides = stride if type(stride) is tuple else (stride, ) * length
     dilations = dilation if type(dilation) is tuple else (dilation, ) * length
+    # `nn.ConvTransposeNd.forward(x, output_size=...)` computes the per-dimension
+    # output_padding as a list, so a tuple check alone would wrap it and the shape
+    # arithmetic below would add a list to an int.
+    output_paddings = output_padding if isinstance(output_padding, (tuple, list)) else (output_padding, ) * length
 
     output_dims = []
     for idx, input_dim in enumerate(input_dims):
-
-        output_dim = (input_dim + 2 * paddings[idx] - (dilations[idx] *
-                                                       (kernel_dims[idx] - 1) + 1)) // strides[idx] + 1
+        # A transposed convolution inverts the forward shape formula, so the output grows
+        # with the stride rather than shrinking by it.
+        output_dim = ((input_dim - 1) * strides[idx] - 2 * paddings[idx] + dilations[idx] * (kernel_dims[idx] - 1) +
+                      output_paddings[idx] + 1)
         output_dims.append(output_dim)
 
-    paddings = padding if type(padding) is tuple else (padding, padding)
-    strides = stride if type(stride) is tuple else (stride, stride)
-    dilations = dilation if type(dilation) is tuple else (dilation, dilation)
-
-    filters_per_channel = out_channels // groups
     conv_per_position_macs = int(_prod(kernel_dims)) * in_channels * filters_per_channel
     active_elements_count = batch_size * int(_prod(input_dims))
     overall_conv_macs = conv_per_position_macs * active_elements_count

tests/unit/profiling/flops_profiler/test_flops_profiler.py

@@ -53,6 +53,62 @@ def forward(self, query, key, value):
     assert profiles == [(65536, 32768)] * 3
 
 
+@pytest.mark.sequential
+@pytest.mark.parametrize("groups", [1, 2, 4])
+@pytest.mark.parametrize("stride, padding, output_padding", [(1, 0, 0), (2, 1, 1)])
+def test_conv_transpose_flops(groups, stride, padding, output_padding):
+    """A transposed convolution scatters every input element across the kernel and adds the
+    bias once per output element, so the two counts follow different shapes."""
+    in_channels, out_channels, kernel_size = 4, 8, 3
+    model = torch.nn.ConvTranspose2d(in_channels,
+                                     out_channels,
+                                     kernel_size,
+                                     stride=stride,
+                                     padding=padding,
+                                     output_padding=output_padding,
+                                     groups=groups)
+    inputs = torch.randn(2, in_channels, 8, 8)
+
+    prof = FlopsProfiler(model)
+    prof.start_profile()
+    outputs = model(inputs)
+    prof.stop_profile()
+    flops, macs = prof.get_total_flops(), prof.get_total_macs()
+    prof.end_profile()
+
+    input_elements = inputs.shape[0] * inputs[0, 0].numel()
+    output_elements = outputs.shape[0] * outputs[0, 0].numel()
+    expected_macs = input_elements * in_channels * (out_channels // groups) * kernel_size * kernel_size
+    expected_bias_flops = out_channels * output_elements
+
+    assert macs == expected_macs
+    assert flops == 2 * expected_macs + expected_bias_flops
+
+
+@pytest.mark.sequential
+def test_conv_transpose_flops_with_output_size():
+    """`output_size=` makes torch compute the per-dimension output_padding itself, and it
+    hands it over as a list rather than a tuple."""
+    in_channels, out_channels, kernel_size = 4, 8, 3
+    model = torch.nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=2, padding=1)
+    inputs = torch.randn(2, in_channels, 8, 8)
+
+    prof = FlopsProfiler(model)
+    prof.start_profile()
+    outputs = model(inputs, output_size=(2, out_channels, 16, 16))
+    prof.stop_profile()
+    flops, macs = prof.get_total_flops(), prof.get_total_macs()
+    prof.end_profile()
+
+    input_elements = inputs.shape[0] * inputs[0, 0].numel()
+    output_elements = outputs.shape[0] * outputs[0, 0].numel()
+    expected_macs = input_elements * in_channels * out_channels * kernel_size * kernel_size
+    expected_bias_flops = out_channels * output_elements
+
+    assert macs == expected_macs
+    assert flops == 2 * expected_macs + expected_bias_flops
+
+
 class LeNet5(torch.nn.Module):
 
     def __init__(self, n_classes):