FLOPs Estimation (Linear/Conv)

The problem, from first principles

02-parameter-counting counts how many numbers a model has, and 03-memory-footprint-estimation counts how many bytes storing them costs — neither says anything about how much compute a forward pass actually takes. Two models can have the same parameter count and wildly different runtime cost (a Linear layer applied once vs. a Conv2d layer's kernel sliding across every spatial position), so a separate estimate is needed: FLOPs, the number of individual floating-point multiply/add operations a forward pass performs.

From theory to code

Implement linear_flops(batch_size, in_features, out_features, bias=True) and conv2d_flops(batch_size, in_channels, out_channels, kernel_size, output_height, output_width, bias=True), each returning the total FLOP count for one forward pass of that layer type.

Constraints

• Every multiply-accumulate (one multiply, one add) counts as 2 FLOPs — the standard convention used by real profiling tools.
• linear_flops: 2 * batch_size * in_features * out_features multiply-add FLOPs, plus batch_size * out_features more if bias=True (one addition per output element).
• conv2d_flops: assumes a square kernel, no groups, no dilation. Each of batch_size * out_channels * output_height * output_width output positions costs in_channels * kernel_size * kernel_size multiply-adds; add one more FLOP per output position if bias=True.
• bias defaults to True in both functions.

Hints

Open one at a time. Each gives away a little more than the last.