Implement and vectorize NumPy Conv2D
Company: Tesla
Role: Machine Learning Engineer
Category: Data Manipulation (SQL/Python)
Difficulty: medium
Interview Round: Technical Screen
Implement a 2D convolution operation from scratch using NumPy only (no TensorFlow or PyTorch). Assume NCHW input shape (N, C_in, H_in, W_in) and weights of shape (C_out, C_in, k_h, k_w); support configurable stride and padding. First provide a clear nested-loop reference implementation; then optimize by vectorizing the computation (e.g., im2col or stride tricks) and discuss time/memory trade-offs.
Overview: This question evaluates understanding of 2D convolution mechanics, multidimensional NumPy array manipulation, and the competency to optimize numerical computations for performance and memory.
You are given an image tensor and a convolution kernel stored in relational tables. The input tensor follows an NCHW layout (batch, channel, height, width), flattened into rows, and the kernel follows a (C_out, C_in, k_h, k_w) layout.
Tables:
- input_tensor: one row per pixel of the input tensor.
- kernel_weights: one row per weight in the convolution kernel.
- conv_config: convolution hyperparameters (stride and padding).
For this data set, there is a single batch (batch_id = 1) and a single configuration (config_id = 1) with stride_h = 1, stride_w = 1, pad_h = 0, pad_w = 0.
Define the 2D convolution as a cross-correlation (no kernel flipping):
output[batch_id, out_channel, out_row, out_col] =
sum over in_channel, k_row, k_col of
input[batch_id, in_channel, out_row * stride_h + k_row - pad_h,
out_col * stride_w + k_col - pad_w]
* kernel[out_channel, in_channel, k_row, k_col]
Assume padding is chosen so that the kernel never goes out of bounds for this sample (i.e., only valid positions are produced). Pixels outside the input tensor are treated as zero (but for this configuration, no such positions are needed).
Write a single SQL query that returns the convolution result as one row per output element, with the following columns:
- batch_id
- out_channel
- out_row (0-based)
- out_col (0-based)
- value (the convolution result at that position)
Your solution should:
- Join input_tensor with kernel_weights and use the stride/padding values from conv_config (config_id = 1).
- Derive the output spatial indices (out_row, out_col) from the input and kernel coordinates.
- Use set-based SQL (joins and aggregations) rather than procedural loops or cursors.
Tables
input_tensor(batch_id INT, in_channel INT, row_idx INT, col_idx INT, value DECIMAL(6,2))
kernel_weights(out_channel INT, in_channel INT, k_row_idx INT, k_col_idx INT, weight DECIMAL(6,2))
conv_config(config_id INT, stride_h INT, stride_w INT, pad_h INT, pad_w INT)
Hints
- Join input_tensor with kernel_weights on in_channel and compute output coordinates from the input and kernel indices.
- Use GROUP BY on the derived (out_row, out_col, batch_id, out_channel) to aggregate the sum of input * weight for each output position.