Implement a Dense Block with 2D Convolutions

#137 · Deep Learning · Hard

Problem

Implement a Dense Block as used in DenseNet architectures. In a Dense Block, each layer receives feature maps from all preceding layers (via concatenation) and passes its output to all subsequent layers. Each layer applies BN-ReLU-Conv.

Solution

import numpy as np

def batch_norm(x, eps=1e-5):
    mean = np.mean(x, axis=(0, 2, 3), keepdims=True)
    var = np.var(x, axis=(0, 2, 3), keepdims=True)
    return (x - mean) / np.sqrt(var + eps)

def relu(x):
    return np.maximum(0, x)

def conv2d(x, kernel, bias, stride=1, pad=1):
    if pad > 0:
        x = np.pad(x, ((0,0),(0,0),(pad,pad),(pad,pad)), mode='constant')
    N, C_in, H, W = x.shape
    C_out, _, kH, kW = kernel.shape
    out_h = (H - kH) // stride + 1
    out_w = (W - kW) // stride + 1
    out = np.zeros((N, C_out, out_h, out_w))
    for n in range(N):
        for co in range(C_out):
            for i in range(out_h):
                for j in range(out_w):
                    hs, ws = i * stride, j * stride
                    out[n, co, i, j] = np.sum(x[n, :, hs:hs+kH, ws:ws+kW] * kernel[co]) + bias[co]
    return out

class DenseLayer:
    def __init__(self, in_channels, growth_rate):
        self.kernel = np.random.randn(growth_rate, in_channels, 3, 3) * np.sqrt(2.0 / (in_channels * 9))
        self.bias = np.zeros(growth_rate)

    def forward(self, x):
        out = batch_norm(x)
        out = relu(out)
        out = conv2d(out, self.kernel, self.bias, stride=1, pad=1)
        return out

class DenseBlock:
    def __init__(self, in_channels, growth_rate, num_layers):
        self.layers = []
        for i in range(num_layers):
            layer_in = in_channels + i * growth_rate
            self.layers.append(DenseLayer(layer_in, growth_rate))

    def forward(self, x):
        features = [x]
        for layer in self.layers:
            concat_input = np.concatenate(features, axis=1)
            new_features = layer.forward(concat_input)
            features.append(new_features)
        return np.concatenate(features, axis=1)

Explanation

A Dense Block consists of multiple dense layers. Each layer receives the concatenation of all previous feature maps.
Each layer applies Batch Normalization, ReLU, and a 3x3 convolution, producing growth_rate new feature maps.
The layer's output is appended to the list of features and concatenated as input to subsequent layers.
The final output concatenates all feature maps: original input plus outputs from every layer.
Total output channels = in_channels + num_layers * growth_rate.

Complexity

Time: O(L^2 growth_rate C H W * k^2) where L = layers, C = cumulative channels, k = kernel size
Space: O(N (in_channels + L growth_rate) H W) for all stored feature maps

← #136 #138 →