Efficient Real-Time Video Colorization on Low-End CPUs via Pruning and Quantization

Skip connections, also known as residual connections, are an essential technique in deep learning that allows neural networks to bypass intermediate layers by adding the input of a layer directly to its output. Formally, for an input x, instead of only learning the transformation f(x), a skip connection computes y = f(x) + x.

This helps preserve information from earlier layers and mitigates the vanishing gradient problem [13], where gradients become too small to update weights in deep networks effectively.

In tasks like video colorization, skip connections are particularly useful in architectures such as U-Net, where they bridge the encoder and decoder layers. The encoder compresses the input image, while the decoder reconstructs it. Skip connections help the decoder access high-resolution features from the encoder, allowing for better preservation of details like edges and textures.

Pruning is a widely used model compression technique that reduces the computational and memory costs of neural networks by removing unnecessary parameters. In convolutional neural networks (CNNs), this typically involves eliminating filters or channels from convolution layers, which reduces the number of operations and speeds up inference without a significant loss in accuracy [5]. Pruning can be either unstructured, where individual weights are removed, or structured, where entire filters or channels are pruned, the latter being more suitable for deployment on hardware due to its regularity.

Magnitude-based pruning is one of the most common techniques, where filters with the smallest weight magnitudes (L1 or L2 norms) are removed [29]. This method works effectively because many neural networks are over-parameterized, meaning that some filters learn redundant or irrelevant features. By pruning those with the least impact on performance, the model becomes smaller and faster while retaining most of its original accuracy.

Quantization is a model compression technique that reduces the precision of a neural network's weights and activations, converting high-precision floating-point values (e.g., 32-bit) into lower precision representations (e.g., 8-bit integers). This significantly reduces the model's memory footprint and computational complexity, making it ideal for deployment in resource-constrained environments, such as edge devices or low-end CPUs.

During quantization, the network maps a range of floating-point values to a smaller set of fixed-point values. Using a representative input sample, the typical numerical range of a neuron is determined and the average of that range is used as an offset. For example, if a neuron is allowed to generate values from − 16 to + 15 using 4 bits but generates values ranging from 2 to 9, the quantization pipeline will take $\frac{2+10}{2}=6$ as the constant offset and use a reduced range of − 4 to + 3 using only 2 bits. The original value can be represented as a sum of the offset and the 2 bit representation, allowing some loss in precision. In practice, quantization reduced the model parameters from 32 bits to 8 bits with minimal impact on performance.

Selecting a training dataset is crucial in achieving robust performance for real-time video colorization. We initially considered the Places365 dataset [44] due to its extensive collection of images from diverse scenes. However, the sheer size of Places365 exceeded the computational resources available for our experiments. To address this challenge and maintain the desired dataset diversity, we opted for a combined approach utilizing two publicly available datasets: VQA-v2 [3] and DAVIS [34]. VQA-v2 offers a rich collection of images covering various real-world scenes, while DAVIS provides high-quality video sequences suitable for training in a video colorization context. Specifically, we employed the Train and Eval splits from the VQA-v2 dataset and the Train split from the DAVIS dataset for model training. We used the separate Eval split from the DAVIS dataset to evaluate the generalizability and performance of our colorization model. This ensures a robust evaluation process that utilizes unseen data during the final assessment.

Following [25], we use a U-Net [35] architecture for image/video colorization. A typical U-Net architecture includes an encoder and a decoder, with multiple skip connections in intermediate layers. Figure 2 shows the overall architecture of the proposed model. The input to the model is the grayscale version of each image, and the ground truth is the RGB image. Intuitively, the model parameters must perform two implicit functions: object identification and color memorization/assignment. Firstly, the model must implicitly identify different objects in the greyscale image. Secondly, it must memorize the color of said objects during training and recall it during inference.

EfficientNet-B7 as U-Net Encoder: Previous research has demonstrated that utilizing a pretrained vision model as the encoder combined with a lightweight decoder achieves an optimal balance between computational efficiency and colorization performance [25]. Due to our focus on real-time video colorization on low-end hardware, we pick EffientNet-B7 [37] as the encoder layer. EfficentNet utilizes depthwise separable convolutions that reduce the computational complexity of convolution. EfficientNet also utilizes skip connections within itself which allows deeper CNN stacks by mitigating the vanishing gradient problem.

Lightweight U-Net Decoder: We use a 9-block decoder, where each block consists of a convolution layer, batch normalization, ReLU activation, and an upsampler. The upsampler uses the nearest neighbor upsampling algorithm and doesn't contain trainable parameters. We experiment with transposed convolution layers but notice no performance difference, corroborated in [25]. From the 82 million parameters in the complete U-Net, only 22% belong to the decoder while the rest belongs to the EfficientNet encoder.

Skip Connections: Skip connections between intermediate layers of the encoder and the decoder are integral architectural choices of the U-Net architecture [35]. Unlike ResNet [7], the purpose of U-Net skip connections is not to counter the vanishing gradient problem [13], rather it facilitates information transfer, especially structural information, from intermediate encoder layer representations to corresponding the decoder layers. Much of the structural information is collapsed when passing through the U-Net bottleneck layer and as such skip connections are vital to the effectiveness of U-Nets for image-to-image tasks such as segmentation and colorization.

Inspired by [25], we used the pixel-wise Weighted Mean Squared Error loss function for our training. Losses incurred on pixels with rarer colors (as computed on the training dataset) are penalized more.

In Equation 1, h, w are the dimensions of the image, Y⁽ⁱ⁾ is the ground truth colored image tensor, $\widehat{Y}_{h, w}^{(i)}$ is the model predicted tensor and ν is the color weighting function. Since keeping track of all 256 × 256 × 256 color combinations is computationally and memory intensive, we bin the colors into ∼ 1 million bins after partitioning the color spectrum into 10 segments. We implement the lookup necessary for ν using an efficient embedding layer ² which is pre-computed and frozen during training.

To reduce color patches and color consistency across frames, we use perceptual loss [17] (See Equation 2). It measures the difference between the features extracted by a pre-trained CNN for the generated colorized frame and the ground truth color frame. Following [33], we use VGG11 [36] to compute perceptual loss. Perceptual loss and by extension, the VGG11 model is used only during training and the VGG11 model is used only for inference and never trained, keeping the pipeline efficient.

In Equation 2, ϕ is the feature tensor from VGG11 for our model's output, and $\widehat{\phi }$ is the ground truth image's feature tensor from VGG11. After pilot experiments, we determined that unweighted addition yields the best colorization and hence, our combined loss is:

Pruning has been widely studied in the context of convolutional neural networks (CNN). The authors of [5] establish a train-prune-retrain framework that we follow in our experiments. Each convolution layer of a CNN has multiple channels or filters that each learn different representations of the image. For example, in a JPEG format image, there are three channels Red, Green, and Blue. Each of the colors is a representation of the complete image and usually contains most of the salient information of the image. The per-channel representations learned by a convolution layer follow the same principle, although are less human-interpretable. Channels Convolution layers near the start of the model often learn to detect edges while upper layers detect entire objects [19].

convolution layers are often intentionally over-parameterized to improve training dynamics [2] and multiple channels converge during training and produce the same output. Since the output of the channels is averaged by the next convolution layer, they may also cancel out each other's contributions. Although redundant parameters in the form of excess channels at each convolution layer make training easier, they do not contribute to downstream performance and rather slow down the inference since redundant computations are being performed.

As such, we choose to prune redundant channels at suitable convolution layers, which is a structured pruning strategy [5] and use magnitude-based pruning [21].

How Skip Connections in CNNs limit Pruning: Algorithm 1 outlines the approach to pruning the trained U-Net. convolution layers are sequentially dependent since the channel dimension of the parameter weights of layer_{i + 1} attend to the outputs of layer_i. In other words, if layer_i has n output channels, layer_{i + 1} must have n input channels. Consequently, if any output channel of layer_i is pruned, the corresponding input channels of layer_{i + 1} must be pruned as well. This is highly challenging in practice due to the skip connections present in modern CNNs, including EfficientNet and U-Net. If both layer_a and layer_b lead to layer_c, pruning the first two layers will mandate pruning twice as many input layers in layer_c. Since modern CNNs have multiple such skip connections attaching to any layer, almost all of the input channels of that layer will be pruned, resulting in a disconnected model. Figure 3 highlights this dilemma.

As such, we opt to prune only layers in our model where there are no skip connections, i.e. the bottleneck layers. Although a majority of the convolution layers cannot be pruned due to skip connections, pruning only the bottleneck layers provide significant performance improvements. As shown in Table 2, we determine target layers from pruning automatically and experimentally find optimal pruning ratios for each category of convolution layer. We did not experiment with tuning the pruning ratio of each individual layer but we expect it to yield further improvements. After pruning, we retrain the pruned model to recover performance.

We use the PyTorch [32] native quantization pipeline for activation-aware quantization of the decoder layers of our trained U-Net. Using sample inputs, the pipeline determines the average range of values that a particular neuron generates and uses the mean of that range as the offset as explained in Section 3.3.

Pruning layers result in the model producing discolored images. This is fixed by fine-tuning the pruned model to adapt to its lost weight, as demonstrated in 4.

Table compares the performance of the baseline model, pruned model, and pruned+quantized model in terms of FPS (frames per second), FPS gain, and model size. The baseline model achieves 4.54 FPS with a size of 330 MB. After pruning, the FPS improves to 6.22, marking a 36.93% increase in speed, while the model size reduces to 237 MB. Further quantization brings the FPS to 6.53, representing a 43.75% speed gain over the baseline, with a model size reduced to 229 MB. This optimization demonstrates both a significant increase in inference speed and a decrease in model size, making it more suitable for deployment on low-end CPUs.

Our model demonstrates promising results in terms of the Color Distribution Consistency (CDC) score. This metric evaluates the consistency of predicted color distributions across the video. As shown in Figure 4, our model achieves a CDC score comparable to state-of-the-art methods. This suggests that the model can predict relatively consistent color distributions within individual frames, even though these predictions may vary in a spatio-temporal manner.

As shown in Figure 4(c), the trained U-Net model tends to produce color patches that appear throughout each video sequence. These patches may vary in size and location across different frames, leading to a lack of visual consistency. The colorization process lacks spatio-temporal coherence, resulting in inconsistencies in color across consecutive frames. This effect manifests as fluctuations in the Peak Signal-to-Noise Ratio (PSNR) plot when plotted against frame number (see Figure 5). The significant variations in PSNR between frames highlight the inconsistency in color prediction across the video sequence.