Skip to content

← Flax step 15 of 100

Hard Primitives

Implement Depthwise-Separable Convolution

Jump to the task

Filter positions, then mix channels

A depthwise separable convolution performs two operations. First, a depthwise filter processes each input channel independently. Then a pointwise filter combines those channel outputs at each image position. This exercise uses one depthwise output per input channel, also called a channel multiplier of one.

There is no activation or normalization between the two operations, and only the final output receives a bias. These details matter: adding a ReLU, a depthwise bias, or extra normalization changes the function and will not match the requested result.

Understand the grouped kernel

The input has shape (H, W, C_in). Add one batch axis before calling jax.lax.conv_general_dilated, giving (1, H, W, C_in).

Use dimension_numbers=("NHWC", "HWIO", "NHWC") for both operations:

Stage Kernel shape Feature groups Output channels
Depthwise (kernel_h, kernel_w, 1, C_in) C_in C_in
Pointwise (1, 1, C_in, features) 1 features

In a grouped convolution, the kernel's input-channel dimension counts channels per group. Each depthwise group sees one input channel, hence the 1 in its kernel shape. The last kernel dimension counts total output channels. Setting the group count alone does not require one output per input; the chosen output-channel count makes that true here. Other depthwise layers can use a larger channel multiplier.

The pointwise operation has no spatial window beyond the current position. It is equivalent to multiplying each position's channel vector by a matrix of shape (C_in, features).

A small channel-routing example

With a one-pixel input [1, 2] and depthwise weights [2, 3], the separate channel results are [2, 6]. A pointwise matrix [[1, -1], [0.5, 2]] turns those into [5, 10]. Adding bias [0.25, -0.5] gives [5.25, 9.5]:

import jax
import jax.numpy as jnp
import numpy as np

x = jnp.array([1., 2.]).reshape(1, 1, 1, 2)
depthwise = jnp.array([2., 3.]).reshape(1, 1, 1, 2)
h = jax.lax.conv_general_dilated(
    x, depthwise, (1, 1), "SAME",
    dimension_numbers=("NHWC", "HWIO", "NHWC"),
    feature_group_count=2,
)
pointwise = jnp.array([[1., -1.], [0.5, 2.]])
y = h[0] @ pointwise + jnp.array([0.25, -0.5])
np.testing.assert_allclose(y.reshape(-1), [5.25, 9.5])

Count parameters without promising a speedup

A regular convolution with bias has kernel_h * kernel_w * C_in * features + features scalar parameters. This layer has kernel_h * kernel_w * C_in + C_in * features + features.

For a 3-by-3 kernel with 64 input and 64 output channels, the weight counts are 36,864 versus 4,672, plus 64 bias values in either case. That is about 7.9 times fewer weights. Actual runtime depends on the shapes, hardware, and implementation, and accuracy depends on the trained model. Neither improvement follows from the count alone.

This factorization cannot represent every regular convolution of the same size. Without intermediate nonlinearities, its equivalent regular kernel is K[h, w, i, o] = depthwise[h, w, 0, i] * pointwise[0, 0, i, o]. For each input channel, all output channels share one spatial filter, scaled by different pointwise weights.

Problem

Implement my_depthwise_separable_forward(seed, x, features, kernel_h, kernel_w):

  1. Define a compact MyDepthwiseSeparable module with integer fields features, kernel_h, and kernel_w.
  2. Create "depthwise_kernel" with shape (kernel_h, kernel_w, 1, x.shape[-1]), initialized by nn.initializers.lecun_normal().
  3. Create "pointwise_kernel" with shape (1, 1, x.shape[-1], features), also using lecun_normal().
  4. Create "bias" with shape (features,), using nn.initializers.zeros.
  5. Add a batch axis. Apply the depthwise convolution with feature_group_count=x.shape[-1], followed by the pointwise convolution with the default group count of one.
  6. Use stride (1, 1), padding "SAME", no dilation, and the layouts above for both calls. Remove the batch axis with h[0], then add bias.
  7. Initialize the module directly on x with jax.random.PRNGKey(int(seed)), apply it with those variables, and flatten the result in row-major order.

Keep the parameter names and creation order above. Do not split the initialization key or add a parent module.

Inputs: x is a float array (H, W, C_in) with positive dimensions. features, kernel_h, and kernel_w are positive integer-valued scalars, possibly floats; convert them to Python integers. seed is an integer-valued scalar. Output: a one-dimensional array of length H * W * features.

SAME padding with stride one preserves height and width. For an even kernel dimension, the extra padding goes on its bottom or right side. Keep all scalar setup arguments static when JIT-compiling the wrapper.