Depthwise Separable Convolution - PyTorch Basics

Depthwise Separable Convolution

What you'll build: the two halves of a depthwise separable convolution: depthwise_conv (each channel filtered by its own kernel) and pointwise_conv (a 1x1 convolution that mixes channels).

Background

A standard convolution filters space and mixes channels in one step. MobileNet and Xception factor it into two cheaper steps: a depthwise convolution that filters each channel spatially on its own, then a pointwise 1x1 convolution that mixes the channels at each pixel. For a K x K kernel the parameter count drops from C_in * C_out * K² to C_in * K² + C_in * C_out: for C_in = C_out = 256, K = 3, from 589,824 to 67,840, about 8.7 times fewer.

Formulas

"Valid" convolution (no padding), stride 1, computed as a cross-correlation (the kernel is not flipped, as in every deep learning library):

depthwise:  out[c, i, j] = sum_{u, v} x[c, i + u, j + v] * kernel[c, u, v]
            output shape (C, H - kH + 1, W - kW + 1)

pointwise:  out[o, i, j] = sum_c kernel[o, c, 0, 0] * x[c, i, j]
            output shape (C_out, H, W)

Inputs

• depthwise_conv(x, kernel): x (C, H, W), kernel (C, kH, kW) (one filter per channel).
• pointwise_conv(x, kernel): x (C_in, H, W), kernel (C_out, C_in, 1, 1) (the layout of a PyTorch Conv2d weight).

Expected Return Format

• depthwise_conv returns (C, H - kH + 1, W - kW + 1).
• pointwise_conv returns (C_out, H, W).

Constraints

• No batch dimension, no padding, no bias, stride 1.
• Loops over output positions are acceptable; a vectorised version uses np.lib.stride_tricks.sliding_window_view and np.einsum.
• PyTorch track: don't use F.conv2d or nn.Conv2d (in PyTorch, depthwise is Conv2d(C, C, k, groups=C)).

Follow-up

• Compute the multiply-adds of both versions for a 56x56x128 input and a 3x3 kernel with 128 output channels.
• Why do MobileNetV2 blocks put the depthwise convolution between two pointwise ones (an inverted residual)?

Input: depthwise_conv(np.ones((1, 3, 3)), np.ones((1, 2, 2)))
Output: array([[[4., 4.], [4., 4.]]])

Input: pointwise_conv(np.array([[[1.0]], [[2.0]]]), np.array([[[[1.0]], [[1.0]]], [[[2.0]], [[-1.0]]]]))
Output: array([[[3.]], [[0.]]])

import numpy as np

def depthwise_conv(x, kernel):
    """
    Depthwise convolution (valid, stride 1, no kernel flip).

    Args:
        x: Input (C, H, W)
        kernel: One filter per channel (C, kH, kW)

    Returns:
        Output (C, H-kH+1, W-kW+1)
    """
    # Your code here
    pass

def pointwise_conv(x, kernel):
    """
    Pointwise (1x1) convolution.

    Args:
        x: Input (C_in, H, W)
        kernel: Filters (C_out, C_in, 1, 1)

    Returns:
        Output (C_out, H, W)
    """
    # Your code here
    pass