Implement 2D convolution forward pass
TeslaMachine Learning Engineer · Technical ScreenLast updated
This question evaluates understanding of 2D convolution mechanics and practical low-level tensor manipulation, including kernel shape handling, padding and stride effects, multi-channel input/output interactions, and optional bias integration.
Problem
Implement the forward pass of a 2D convolution (conv2d) from scratch (no deep learning libraries).
You are given:
- Input tensor
xwith shape (N, C_in, H, W) (NCHW layout) - Filter weights
wwith shape (C_out, C_in, K_h, K_w) - Optional bias
bwith shape (C_out,) (may beNone) - Integer stride
s_h, s_w - Integer padding
p_h, p_w(zero-padding applied to height/width)
Compute the output tensor y with shape (N, C_out, H_out, W_out) where:
[ H_{out} = \left\lfloor \frac{H + 2p_h - K_h}{s_h} \right\rfloor + 1, \quad W_{out} = \left\lfloor \frac{W + 2p_w - K_w}{s_w} \right\rfloor + 1 ]
For each output element:
[ y[n, c_{out}, i, j] = \sum_{c_{in}=0}^{C_{in}-1} \sum_{u=0}^{K_h-1} \sum_{v=0}^{K_w-1} ; x_{pad}[n, c_{in}, i\cdot s_h + u, j\cdot s_w + v] \cdot w[c_{out}, c_{in}, u, v] + b[c_{out}] ]
where x_pad is x padded with zeros by (p_h, p_w).
Requirements
- Return
yas a dense numeric array/tensor. - Do not use existing convolution operators.
- Handle edge cases such as
b is None, non-square kernels, and different strides/padding.
Constraints (typical for an interview unit test)
- Sizes are small enough that an (O(N\cdot C_{out}\cdot C_{in}\cdot H_{out}\cdot W_{out}\cdot K_h\cdot K_w)) implementation passes.
- Inputs are floating point numbers.
Example (shape check)
If x is (1, 3, 32, 32), w is (8, 3, 3, 3), stride (1,1), padding (1,1), then output shape is (1, 8, 32, 32).
Solution walkthrough is for members
Get the full walkthrough with reference code and complexity analysis, plus Kitty Coach at 300 messages a day. Hints, the code runner and community answers stay free.