Structured Knowledge Distillation for Dense Prediction

Introduction

Dense prediction is a family of fundamental problems in computer vision, which learns a mapping from input images to complex output structures, including semantic segmentation, depth estimation and object detection, among many others. One needs to assign category labels or regress specific values for each pixel given an input image to form the structured outputs. In general these tasks are significantly more challenging to solve than image-level prediction problems, thus often requiring networks with large capacity in order to achieve satisfactory accuracy. On the other hand, compact models are desirable for enabling computing on edge devices with limited computation resources.

Deep neural networks have been the dominant solutions since the invention of fully-convolutional neural networks (FCNs) . Subsequent approaches, e.g., DeepLab , PSPNet , RefineNet , and FCOS follow the design of FCNs to optimize energy-based objective functions related to different tasks, having achieved significant improvement in accuracy, often with cumbersome models and expensive computation.

Recently, design of neural networks with compact model sizes, light computation cost and high performance, has attracted much attention due to the need of applications on mobile devices. Most current efforts have been devoted to designing lightweight networks specifically for dense prediction tasks or borrowing the design from classification networks, e.g., ENet , ESPNet and ICNet for semantic segmentation, YOLO and SSD for object detection, and FastDepth for depth estimation. Strategies such as pruning , knowledge distillation are applied to helping the training of compact networks by making use of cumbersome networks.

Knowledge distillation has proven effective in training compact models for classification tasks . As for dense prediction tasks, most previous works directly apply distillation at each pixel separately to transfer the class probability from the cumbersome network (teacher) to the compact network (student); or to extract more discriminative feature embeddings for the compact network. Note that, such a pixel-wise distillation scheme neglects the important structure information.

Considering the characteristic of dense prediction problem, here we present structured knowledge distillation and transfer the structure information with two schemes, pair-wise distillation and holistic distillation. The pair-wise distillation scheme is motivated by the widely-studied pair-wise Markov random field framework for enforcing spatial labeling consistency. The goal is to align a static affinity graph which is computed to capture both short and long range structure information among different locations from the compact network and the teacher network.

The holistic distillation scheme aims to align higher-order consistencies, which are not characterized in the pixel-wise and pair-wise distillation, between output structures produced from the compact network and the teacher network. We adopt the adversarial training scheme, and a fully convolutional network, a.k.a. the discriminator, considering both the input image and the output structures to produce a holistic embedding which represents the quality of the structure. The compact network is encouraged to generate structures with similar embeddings as the teacher network. We encode the structure knowledge into the weights of discriminators.

To this end, we optimize an objective function that combines a conventional task loss with the distillation terms. The main contributions of this paper are as follows.

We study the knowledge distillation strategy for training accurate compact networks for dense prediction.

We present two structured knowledge distillation schemes, pair-wise distillation and holistic distillation, enforcing pair-wise and high-order consistency between the outputs of the compact and teacher networks.

We demonstrate the effectiveness of our approach by improving recent state-of-the-art compact networks on three different dense prediction tasks: semantic segmentation, depth estimation and object detection. Taking semantic segmentation as an example, the performance gain is illustrated in Figure 1.

Semantic segmentation. Semantic segmentation is a pixel classification problem, which requires an semantic understanding of the whole scene. Deep convolutional neural networks have been the dominant solution to semantic segmentation since the pioneering work, fully-convolutional network . Various schemes have been developed for improving the network capability and accordingly the segmentation performance. For example, stronger backbone networks such as ResNet , DenseNet and HRNet [SunXLW19, hrnet], have shown improved segmentation performance. Retaining the spatial resolution through dilated convolutions or multi-path refine networks leads to significant performance gain. Exploiting multi-scale context using dilated convolutions , or pyramid pooling modules as in PSPNet , also benefits the segmentation. Lin et al. combine deep models with structured output learning for semantic segmentation.

Recently, highly efficient segmentation networks have been attracting increasingly more interests due to the need of real-time and mobile applications. Most works focus on lightweight network design by accelerating the convolution operations with techniques such as factorization techniques. ENet , inspired by , integrates several acceleration factors, including multi-branch modules, early feature map resolution down-sampling, small decoder size, filter tensor factorization, and so on. SQ adopts the SqueezeNet fire modules and parallel dilated convolution layers for efficient segmentation. ESPNet proposes an efficient spatial pyramid, which is based on filter factorization techniques: point-wise convolutions and spatial pyramid of dilated convolutions, to replace the standard convolution. The efficient classification networks such as MobileNet and ShuffleNet are also applied to accelerate segmentation. In addition, ICNet (image cascade network) exploits the efficiency of processing low-resolution images and high inference quality of high-resolution ones, achieving a trade-off between efficiency and accuracy.

Depth estimation. Depth estimation from a monocular image is essentially an ill-posed problem, which requires an expressive model with large capacity. Previous works depend on hand-crafted features . Since Eigen et al. propose to use deep learning to predict depth maps, following works benefit from the increasing capacity of deep models and achieve good results. Besides, Fei propose a semantically informed geometric loss while Yin et al. introduce a virtual normal loss to exploit geometric information. As in semantic segmentation, some works try to replace the encoder with efficient backbones to decrease the computational cost, but often suffer from the limitation caused by the capacity of a compact network. The close work to ours applies pruning to train the compact depth network. However, here we focus on the structured knowledge distillation.

Object detection. Object detection is a fundamental task in computer vision, in which one needs to regress the location of bounding boxes and predict a category label for each instance of interest in an image. Early works achieve good performance by first predicting proposals and then refining the localization of bounding boxes. Effort is also spent on improving detection efficiency with methods such as Yolo , and SSD . They use a one-stage method and design light-weight network structures. RetinaNet solves the problem of unbalance samples to some extent by proposing the focal loss, which makes the results of one-stage methods comparable to two-stage ones. Most of the above detectors rely on a set of pre-defined anchor boxes, which decreases the training samples and makes the detection network sensitive to hyper parameters. Recently, anchor free methods show promises, e.g., FCOS . FCOS employs a fully convolutional framework, and predicts bounding box based on every pixels same as in semantic segmentation, which solves the object detection task as a dense prediction problem. In this work, we apply the structured knowledge distillation method with the FCOS framework, as it is simple and achieves promising performance.

Knowledge distillation. Knowledge distillation transfers knowledge from a cumbersome model to a compact model so as to improve the performance of compact networks. It has been applied to image classification by using the class probabilities produced from the cumbersome model as “soft targets” for training the compact model or transferring the intermediate feature maps .

The MIMIC method distills a compact object detection network by making use of a two-stage Faster-RCNN . They align the feature map at the pixel level and do not make use of the structure information among pixels.

The recent and independent work, which applies distillation to semantic segmentation , is related to our approach. It mainly distills the class probabilities for each pixel separately (as our pixel-wise distillation) and center-surrounding differences of labels for each local patch (termed as a local relation in ). In contrast, we focus on distilling structured knowledge: pair-wise distillation, which transfers the relation among different locations by building a static affinity graph, and holistic distillation, which transfers the holistic knowledge that captures high-order information. Thus the work of may be seen as a special case of the pair-wise distillation.

This work here is a substantial extension of our previous work appeared in . The main difference compared with is threefold. 1) We extend the pair-wise distillation to a more general case in Section 2.1 and build a static graph with nodes and connections. We explore the influence of the graph size, and find out that it is important to keep a fully connected graph. 2) We provide more explanations and ablations on the adversarial training for holistic distillation. 3) We also extend our method to depth estimation and object detection with two recent strong baselines and , by replacing the backbone with MobileNetV2 , and further improve their performance.

Adversarial learning. Generative adversarial networks (GANs) have been widely studied in text generation and image synthesis . The conditional version is successfully applied to image-to-image translation, such as style transfer and image coloring .

The idea of adversarial learning is also employed in pose estimation , encouraging the human pose estimation result not to be distinguished from the ground-truth; and semantic segmentation , encouraging the estimated segmentation map not to be distinguished from the ground-truth map. One challenge in is the mismatch between the generator’s continuous output and the discrete true labels, making the discriminator in GAN be of very limited success. Different from , in our approach, the employed GAN does not face this issue as the ground truth for the discriminator is the teacher network’s logits, which are real valued. We use adversarial learning to encourage the alignment between the output maps produced from the cumbersome network and the compact network. However, in the depth prediction task, the ground truth maps are not discrete labels. In , the authors use the ground truth maps as the real samples. Different from theirs, our distillation methods attempt to align the output of the cumbersome network and that of the compact network. The task loss calculated with ground truth is optional. When the labelled data is limited, given a well-trained teacher, our method can be applied to unlabelled data and may further improve the accuracy.

Approach

In this section, we first introduce the structured knowledge distillation method for semantic segmentation, a task of assigning a category label to each pixel in the image from $C$ categories. A segmentation network takes a RGB image $\mathbf{I}$ of size $W\times H\times 3$ as the input; then it computes a feature map $\mathbf{F}$ of size $W^{\prime}\times H^{\prime}\times N$, where $N$ is the number of channels. Then, a classifier is applied to compute the segmentation map $\mathbf{Q}$ of size $W^{\prime}\times H^{\prime}\times C$ from $\mathbf{F}$, which is upsampled to the spatial size $W\times H$ of the input image to obtain the segmentation results. We extend our method to other two dense prediction tasks: depth estimation and object detection.

Pixel-wise distillation. We apply the knowledge distillation strategy to transfer the knowledge of the large teacher segmentation network $\mathsf{T}$ to a compact segmentation network $\mathsf{S}$ for better training the compact segmentation network. We view the segmentation problem as a collection of separate pixel labeling problems, and directly use knowledge distillation to align the class probability of each pixel produced from the compact network. We follow and use the class probabilities produced from the teacher model as soft targets for training the compact network.

where $\mathbf{q}^{s}_{i}$ represents the class probabilities of the $i$th pixel produced from the compact network $\mathsf{S}$. $\mathbf{q}^{t}_{i}$ represents the class probabilities of the $i$th pixel produced from the cumbersome network $\mathsf{T}$. $\operatorname{KL}(\cdot)$ is the Kullback-Leibler divergence between two probabilities, and $\mathcal{R}=\{1,2,\dots,W^{\prime}\times H^{\prime}\}$ denotes all the pixels.

In addition to above straightforward pixel-wise distillation, we present two structured knowledge distillation schemes—pair-wise distillation and holistic distillation—to transfer structured knowledge from the teacher network to the compact network. The pipeline is illustrated in Figure 2.

Pair-wise distillation. Inspired by the pair-wise Markov random field framework that is widely adopted for improving spatial labeling contiguity, we propose to transfer the pair-wise relations, specifically pair-wise similarities in our approach, among spatial locations.

We build an affinity graph to denote the spatial pair-wise relations, in which, the nodes represent different spatial locations and the connection between two nodes represents the similarity. We denote the connection range $\alpha$ and the granularity $\beta$ of each node to control the size of the static affinity graph. For each node, we only consider similarities with top-$\alpha$ near nodes according to spatial distance (here we use the Chebyshev distance) and aggregate $\beta$ pixels in a spatial local patch to represent the feature of this node as illustrate in Figure 3. Here for a $W^{\prime}\times H^{\prime}\times C$ feature map, $W^{\prime}\times H^{\prime}$ is the spatial resolution. With the granularity $\beta$ and the connection range $\alpha$, the affinity graph contains $\frac{W^{\prime}\times H^{\prime}}{\beta}$ nodes with $\frac{W^{\prime}\times H^{\prime}}{\beta}\times\alpha$ connections.

Let $a^{t}_{ij}$ and $a^{s}_{ij}$ denote the similarity between the $i$th node and the $j$th node produced from the teacher network $\mathsf{T}$ and the compact network $\mathsf{S}$, respectively. We adopt the squared difference to formulate the pair-wise similarity distillation loss,

where $\mathcal{R^{\prime}}=\{1,2,\dots,\frac{W^{\prime}\times H^{\prime}}{\beta}\}$ denotes all the nodes. In our implementation, we use average pooling to aggregate $\beta\times C$ features in one node to be $1\times C$, and the similarity between two nodes is computed from the aggregated features $\mathbf{f}_{i}$ and $\mathbf{f}_{j}$ as

Holistic distillation. We align high-order relations between the segmentation maps produced from the teacher and compact networks. The holistic embeddings of the segmentation maps are computed as the representations.

We use conditional generative adversarial learning for formulating the holistic distillation problem. The compact net is viewed as a generator conditioned on the input RGB image $\mathbf{I}$, and the predicted segmentation map $\mathbf{Q}^{s}$ is seen as a fake sample. The segmentation map predicted by the teacher ( $\mathbf{Q}^{t}$) is the real sample. We expect that $\mathbf{Q}^{s}$ is similar to $\mathbf{Q}^{t}$. Generative adversarial networks (GANs) usually suffer from the unstable gradient in training the generator due to the non-smoothness of the discontinuous Kullback–Leibler (KL) or Jensen-Shannon (JS) divergence when the two distributions do not overlap. The Wasserstein distance provides a smooth measure of the difference between two distributions. The Wasserstein distance is defined as the minimum cost to converge the model distribution $p_{s}(\mathbf{Q}^{s})$ to the real distribution $p_{t}(\mathbf{Q}^{t})$. This can be written as follows:

The segmentation map and the RGB image are concatenated as the input of the embedding network $\mathsf{D}$. $\mathsf{D}$ is a fully convolutional neural network with five convolution blocks. Each convolution block has ReLU and BN layers, except for the final output convolution block. Two self-attention modules are inserted between the final three blocks to capture the structure information . We insert a batch normalization immediately after the input so as to normalize the scale difference between the segmentation map and RGB channels.

Such a discriminator is able to produce a holistic embedding representing how well the input image and the segmentation map match. We further add a pooling layer to pool the holistic embedding into a score. As we employ the Wasserstein distance in the adversarial training, the discriminator is trained to produce a higher score w.r.t. the output segmentation map from the teacher net and produce lower scores w.r.t. the ones from the student. In this process, we encode the knowledge of evaluating the quality of a segmentation map into the discriminator. The student is trained with the regularization of achieving a higher score under the evaluation of the discriminator, thus improving the performance of the student.

2 Optimization

where $\lambda_{1}$ and $\lambda_{2}$ are set to $10$ and $0.1$, making these loss value ranges comparable. We minimize the objective function with respect to the parameters of the compact segmentation network $\mathsf{S}$, while maximizing it w.r.t. the parameters of the discriminator $\mathsf{D}$, which is implemented by iterating the following two steps:

Train the compact segmentation network $\mathsf{S}$. Given the discriminator network, the goal is to minimize the multi-class cross-entropy loss and the distillation losses relevant to the compact segmentation network:

3 Extension to Other Dense Prediction Tasks

Dense prediction learns a mapping from an input RGB image $\mathbf{I}$ of size $W\times H\times 3$ to a per-pixel output $\mathbf{Q}$ of size $W\times H\times C$. In semantic segmentation, the output has $C$ channels which is the number of semantic classes.

For the object detection task, for each pixel, we predict the $c^{*}$ classes, as well as a 4D vector $t^{*}=(l,t,r,b)$ representing the location of a bounding box. We follow FCOS , and combine them with distillation terms as regularization.

The depth estimation task can be solved as a classification task, as the continuous depth values can be divided into $C$ discrete categories . For inference, we apply a soft weighted sum as in . The pair-wise distillation can be directly applied to the intermediate feature maps. The holistic distillation uses the depth map as input. We can use the ground truth as the real samples of GAN in the depth estimation task, because it is a continuous map. However, in order to apply our method to unlabelled data, we still use depth maps from the teacher as our real samples.

Experiments

In this section, we first apply our method to semantic segmentation to empirically verify the effectiveness of structured knowledge distillation. We discuss and explore how structured knowledge distillation works.

Structured knowledge distillation can be applied to other structured output prediction tasks with FCN frameworks. In Section 3.3 and Section 3.2, we apply our distillation method to strong baselines in object detection and depth estimation tasks with minimum modifications.

Network structures. We use the segmentation architecture of PSPNet with a ResNet$101$ as the teacher network $\mathsf{T}$.

We study recent compact networks, and employ several different architectures to verify the effectiveness of the distillation framework. We first use ResNet$18$ as a basic student network and conduct ablation studies. Then, we employ the MobileNetV$2$Plus , which is based on a pretrained MobileNetV$2$ model on the ImageNet dataset. We also test ESPNet-C and ESPNet models, which are also lightweight models.

Training setup. Unless otherwise specified, segmentation networks here using stochastic gradient descent (SGD) with the momentum ($0.9$) and the weight decay ($0.0005$) for $40000$ iterations. The learning rate is initialized to be $0.01$ and is multiplied by $(1-\frac{iter}{max-iter})^{0.9}$. We random crop the images into the resolution of $512\times 512$ pixels as the training input. Standard data augmentation is applied during training, such as random scaling (from $0.5$ to $2.1$) and random flipping. We follow the settings in to reproduce the results of ESPNet and ESPNet-C, and train the compact networks with our distillation framework.

1.2 Datasets

Cityscapes. The Cityscapes dataset is collected for urban scene understanding and contains $30$ classes with only $19$ classes used for evaluation. The dataset contains $5,000$ high quality finely annotated images and $20,000$ coarsely annotated images. The finely annotated images are divided into $2,975$, $500$, $1,525$ images for training, validation and testing. We only use the finely annotated dataset in our experiments.

CamVid. The CamVid dataset contains $367$ training and $233$ testing images. We evaluate the performance over $11$ different classes such as building, tree, sky, car, road, etc.

ADE$20$K. The ADE$20$K dataset contains $150$ classes of diverse scenes. The dataset is divided into $20K$/$2K$/$3K$ images for training, validation and testing.

1.3 Evaluation Metrics

We use the following metrics to evaluate the segmentation accuracy, model size and the efficiency.

The Intersection over Union (IoU) score is calculated as the ratio of interval and union between the ground-truth mask and the predicted segmentation mask for each class. We use the mean IoU of all classes (mIoU) to study the effectiveness of distillation. We also report the class IoU to study the effect of distillation on different classes. Pixel accuracy is the ratio of the pixels with the correct semantic labels to the overall pixels.

The model size is represented by the number of network parameters. and the complexity is evaluated by the sum of floating point operations (FLOPs) in one forward on a fixed input size.

1.4 Ablation Study

The effectiveness of distillations. We examine the role of each component of our distillation system. The experiments are conducted on ResNet$18$ with its variant ResNet$18$ ($0.5$) representing a width-halved version of ResNet$18$ on the Cityscapes dataset. Table I reports the results of different settings for the student net, which are the average results of the final epoch from three runs.

From Table I, we can see that distillation can improve the performance of the student network, and structure distillation helps the student learn better. With the three distillation terms, the improvements for ResNet$18$ ($0.5$), ResNet$18$ ($1.0$) and ResNet$18$ ($1.0$) initialized with pretrained weights on the ImageNet dataset are $6.26\%$, $5.74\%$ and $2.9\%$, respectively, which indicates that the effect of distillation is more pronounced for the smaller student network and networks without initialization with the weight pretrained from the ImageNet. Such an initialization is also able to transfer the knowledge from another source (ImageNet). The best mIoU of the holistic distillation for ResNet$18$ ($0.5$) reaches $62.7\%$ on the validation set.

On the other hand, each distillation scheme leads to higher mIoU scores. This implies that the three distillation schemes make complementary contributions for better training the compact network.

The affinity graph in pair-wise distillation. In this section, we discuss the impact of the connection range $\alpha$ and the granularity of each node $\beta$ in building the affinity graph. To calculate the pair-wise similarity among each pixel in the feature map will form the fully connected affinity graph, with the price of high computational complexity. We fix the node to be one pixel, and vary the connection range $\alpha$ from the fully connected graph to local sub-graph. Then, we keep the connection range $\alpha$ to be fully connected, and use a local patch to denote each node to change the granularity $\beta$ from fine to coarse. The result are shown in Table II. The results of different settings for the pair-wise distillation are the average results from three runs. We employ a ResNet$18$ ($1.0$) with the weight pretrained from ImageNet as the student network. All the experiments are performed with both pixel-wise distillation and pair-wise distillation, but the sizes of the affinity graph in pair-wise distillation vary.

From Table II, we can see that increasing the connection range can help improve the distillation performance. With the fully connected graph, the student can achieve around $71.37\%$ mIoU. It appears that the best $\beta$ is $2\times 2$, which is slightly better than the finest affinity graph, but the connections are significantly decreased. Using a small local patch to denote a node and calculate the affinity graph may form a more stable correlation between different locations. One can choose to use the local patch to decrease the number of the nodes, instead of decreasing the connection range for a better trade-off between efficiency and accuracy.

To include more structure information, we fuse pair-wise distillation items with different affinity graphs. Three pair-wise fusion strategies are introduced: $\alpha$ fusion, $\beta$ fusion and feature level fusion. The details are shown in Table II. We can see that combining more affinity graphs may improve the performance, but also introduces extra computational cost during training. Therefore, we apply only one pair-wise distillation item in our methods.

Adversarial training in holistic distillation. In this section, we illustrate that GAN is able to encode the holistic knowledge. Details of the discriminator are described in Section 3. The capability of discriminator would affect the adversarial training, and we conduct experiments to discuss the impact of the discriminator’s architecture. The results are shown in Table III. We use $AnLm$ to represent the architecture of the discriminator with $n$ self-attention layers and $m$ convolution blocks with BN layers. The detailed structures can be seen in Figure 4, and the red arrows represent self-attention layers.

From Table III, we can see adding self-attention layers can improve mIoU, and adding more self-attention layers does not change the results much. We add $2$ self-attention blocks considering the performance, stability, and computational cost in our discriminator. With the same self-attention layer, a deeper discriminator can help the adversarial training.

To verify the effectiveness of adversarial training, we further explore the capability of three typical discriminators: the shallowest one (D_Shallow, i.e., A2L2), the one without attention layer (D_no_attention, i.e., A0L4) and ours (D_Ours, i.e., A2L4). The IoUs for different classes are listed in Table IV. It is clear that self-attention layers can help the discriminator better capture the structure, thus the accuracy of the students with the structure objects is improved.

In the adversarial training, the student, a.k.a. the generator, tries to learn the distribution of the real samples (output of the teacher). We apply the Wasserstein distance to transfer the distance between two distributions into a more intuitive score, and the score are highly relevant to the quality of the segmentation maps. We use a well-trained discriminator $\mathsf{D}$ (A2L4) to evaluate the score of a segmentation map. For each image, we feed five segmentation maps, output by the teacher net, the student net w/o holistic distillation, and the student nets w/ holistic distillation under three different discriminator architectures (listed in Table IV) into the discriminator $\mathsf{D}$, and compare the distribution of embedding scores. We evaluate on the validation set and calculate the average score difference between different student nets and the teacher net, the results are shown in Table V. With holistic distillation, the segmentation maps produced from student net can achieve a similar score to the teacher, indicating that GAN helps distill the holistic structure knowledge.

We also draw a histogram to show score distributions of segmentation maps across the validation set in Figure 5. The well-trained $\mathsf{D}$ can assign a higher score to high quality segmentation maps, and the three student nets with the holistic distillation can generate segmentation maps with higher scores and better quality. Adding self-attention layers and more convolution blocks help the student net to imitate the distribution of the teacher net, and attain better performance.

Feature and local pair-wise distillation. We compare a few variants of the pair-wise distillation:

Feature distillation by MIMIC : We follow to align the features of each pixel between $\mathsf{T}$ and $\mathsf{S}$ through a $1\times 1$ convolution layer to match the dimension of the feature

Feature distillation by attention transfer : We aggregate the response maps into an attention map (single channel), and then transfer the attention map from the teacher to the student.

Local pair-wise distillation : This method can be seen as a special case of our pair-wise distillation, which only cover a small sub-graph ($8$-neighborhood pixels for each node).

We replace our pair-wise distillation by the above three distillation schemes to verify the effectiveness of our pair-wise distillation. From Table VI, we can see that our pair-wise distillation method outperforms all the other distillation methods. The superiority over feature distillation schemes (MIMIC and attention transfer , which transfers the knowledge for each pixel separately) may be due the fact that we transfer the structured knowledge other than aligning the feature for each individual pixel. The superiority over the local pair-wise distillation shows the effectiveness of our fully connected pare-wise distillation, which is able to transfer the overall structure information other than a local boundary information .

1.5 Segmentation Results

Cityscapes. We apply our structure distillation method to several compact networks: MobileNetV2Plus which is based on a MobileNetV2 model, ESPNet-C and ESPNet which are carefully designed for mobile applications. Table VII presents the segmentation accuracy, the model complexity and the model size. FLOPsThe FLOPs is calculated with the pytorch version implementation . is calculated on the resolution of $512\times 1024$ pixels to evaluate the complexity. #Params is the number of network parameters. We can see that our distillation approach can improve the results over $5$ compact networks: ESPNet-C and ESPNet , ResNet$18$ ($0.5$), ResNet$18$ ($1.0$), and MobileNetV2Plus . For the networks without pre-training, such as ResNet$18$ ($0.5$) and ESPNet-C, the improvements are very significant with $7.3\%$ and $6.6\%$, respectively. Compared with MD (Enhanced) that uses the pixel-wise and local pair-wise distillation schemes over MobileNet, our approach with the similar network MobileNetV2Plus achieves higher segmentation quality ($74.5$ vs. $71.9$ on the validation set) with a little higher computation complexity and much smaller model size.

Figure 7 shows the IoU scores for each class over MobileNetV2Plus. Both the pixel-wise and structured distillation schemes improve the performance, especially for the categories with low IoU scores. In particular, the structured distillation (pair-wise and holistic) has significant improvement for structured objects, e.g., $17.23\%$ improvement for Bus and $10.03\%$ for Truck. The qualitative segmentation results in Figure 8 visually demonstrate the effectiveness of our structured distillation for structured objects, such as trucks, buses, persons, and traffic signs.

CamVid. Table VIII shows the performance of the student networks w/o and w/ our distillation schemes and state-of-the-art results. We train and evaluate the student networks w/ and w/o distillation at the resolution $480\times 360$ following the setting of ENet. Again we can see that the distillation scheme improves the performance. Figure 9 shows some samples on the CamVid test set w/o and w/ the distillation produced from ESPNet.

We also conduct an experiment by using an extra unlabeled dataset, which contains $2000$ unlabeled street scene images collected from the Cityscapes dataset, to show that the distillation schemes can transfer the knowledge of the unlabeled images. The experiments are done with ESPNet and ESPNet-C. The loss function is almost the same except that there is no cross-entropy loss over the unlabeled dataset. The results are shown in Figure 10. We can see that our distillation method with the extra unlabeled data can significantly improve mIoU of ESPNet-C and ESPNet for $13.5\%$ and $12.6\%$.

ADE$20$K. The ADE$20$K dataset is a very challenging dataset and contains $150$ categories of objects. Note that the number of pixels belonging to difference categories in this dataset is very imbalanced.

We report the results for ResNet$18$ and the MobileNetV$2$ which are trained with the initial weights pretrained on the ImageNet dataset, and ESPNet which is trained from scratch in Table IX. We follow the same training scheme in . All results are tested on single scale. For ESPNet, with our distillation, we can see that the mIoU score is improved by $3.78\%$, and it achieves a higher accuracy with fewer parameters compared to SegNet. For ResNet$18$ and MobileNetV$2$, after the distillation, we achieve $2.73\%$ improvement over the one without distillation reported in .

2 Depth Estimation

Network structures. We use the same model described in with the ResNext$101$ backbone as our teacher model, and replace the backbone with MobileNetV$2$ as the compact model.

Training details. We train the student net using the crop size $385\times 385$ by mini-batch stochastic gradient descent (SGD) with batchsize of $12$. The initialized learning rate is $0.001$ and is multiplied by $(1-\frac{iter}{max-iter})^{0.9}$. For both w/ and w/o distillation methods, the training epoch is $200$.

2.2 Dataset

NYUD-V2. The NYUD-V2 dataset contains $1449$ annotated indoor images, in which $795$ images are for training and others are for testing. The image size is $640\times 480$. Some methods have sampled more images from the video sequence of NYUD-V2 to form a Large-NYUD-V2 to further improve the performance. Following , we conduct ablation studies on the small dataset and also apply the distillation method on current state-of-the-art real-time depth models trained with Large-NYUD-V2 to verify the effectiveness of the structured knowledge distillation.

2.3 Evaluation Metrics

We follow previous methods to evaluate the performance of monocular depth estimation quantitatively based on following metrics: mean absolute relative error (rel), mean $\log_{10}$ error ($\log_{10}$), root mean squared error (rms) , and the accuracy under threshold ($\delta_{i}<1.25^{i},i=1,2,3$).

2.4 Results

Ablation studies. We compare the pixel-wise distillation and the structured knowledge distillation in this section. In the dense classification problem, e.g., semantic segmentation, the output logits of the teacher is a soft distribution of all the classes, which contain the relations among different classes. Therefore, directly transfer the logits from teacher models from the compact ones at the pixel level can help improve the performance. Different from semantic segmentation, as the depth map are real values, the output of the teacher is often not as accurate as ground truth labels. In the experiments, we find that adding pixel-level distillation hardly improves the accuracy in the depth estimation task. Thus, we only use the structured knowledge distillation in depth estimation task.

To verify that the distillation method can further improve the accuracy with unlabeled data, we use $30K$ image sampled from the video sequence of NYUD-V2 without the depth map. The results are shown in Table X. We can see that the structured knowledge distillation performs better than pixel-wise distillation, and adding extra unlabelled data can further improve the accuracy.

Comparison with state-of-the-art. We apply the distillation method to a few state-of-the-art lightweight models for depth estimation. Following , we train the student net on Large-NYUD-V2 with the same constraints in as our baseline, and achieve $13.5$ in the metric ‘rel’. Following the same training setups, with the structured knowledge distillation terms, we further improve the strong baseline, and achieve a relative error (rel) of $13.0$. In Table XI, we list the model parameters and accuracy of a few state-of-the-art large models along with some real-time models, indicating that the structured knowledge distillation works on a strong baseline.

3 Object Detection

Network structures. We experiment with the recent one-stage architecture FCOS , using the backbone ResNeXt-$32$x$8$d-$101$-FPN as the teacher network. The channel in the detector towers is set to $256$. It is a simple anchor-free model, but can achieve comparable performance with state-of-the-art two-stage detection methods.

We choose two different models based on the MobileNetV$2$ backbone: c$128$-MNV$2$ and c$256$-MNV$2$ released by FCOS as our student nets, where $c$ represents the channel in the detector towers. We apply the distillation loss on all the output levels of the feature pyramid network.

Training setup. We follow the training schedule in FCOS . For ablation studies, all the teacher, the student w/ and w/o distillation are trained with stochastic gradient descent (SGD) for $90$K iterations with the initial learning rate being $0.01$ and a mini batch of $16$ images. The learning rate is reduced by a factor of $10$ at iteration $60$K and $80$K, respectively. Weight decay and momentum are set to be $0.0001$ and $0.9$, respectively. To compare with other state-of-the-art real-time detectors, we double the training iterations and the batch size, and the distillation method can further improve the results on the strong baselines.

3.2 Dataset

COCO. Microsoft Common Objects in Context (COCO) is a large-scale detection benchmark in object detection. There are $115K$ images for training and $5K$ images for validation. We evaluate the ablation results on the validation set, and we also submit the final results to the test-dev of COCO.

3.3 Evaluation Metrics

Average precision (AP) computes the average precision value for recall value over to $1$. The mAP is the averaged AP over multiple Intersection over Union (IoU) values, from $0.5$ to $0.95$ with a step of $0.05$. We also report AP$50$ and AP$75$ represents for the AP with a single IoU of $0.5$ and $0.75$, respectively. APs, APm and APl are AP across different scales for small, medium and large objects.

3.4 Results

Comparison with different distillation methods. To demonstrate the effectiveness of the structured knowledge distillation, we compare the pair-wise distillation method with the previous MIMIC method, which aligns the feature map on pixel-level. We use the c$256$-MNV$2$ as the student net and the results are shown in Table XII. By adding the pixel-wise MIMIC distillation method, the detector can be improved by $0.4\%$ in mAP. Our structured knowledge distillation method can improve by $0.9\%$ in mAP. Under all evaluation metrics, the structured knowledge distillation method performs better than MIMIC. By combing the structured knowledge distillation with the pixel-wise distillation, the results can be further improved to $32.1\%$ mAP. Comparing to the baseline method without distillation, the improvement of AP$75$, APs and APl are more sound, indicating the effectiveness of the distillation method.

We show some detection results in Figure 12. One can see that the detector trained with our distillation method can detect more small objects such as ‘person’ and ‘bird’.

Results of different student nets. We follow the same training steps ($90$K) and batch size ($32$) as in FCOS and apply the distillation method on two different structures: C$256$-MV$2$ and C$128$-MV$2$. The results of w/ and w/o distillation are shown in Table XIII. By applying the structured knowledge distillation combine with pixel-wise distillation, the mAP of C$128$-MV$2$ and and C$256$-MV$2$ are improved by $0.9$ and $0.8$, respectively.

Results on the test-dev. The original mAP on the validation set of C$128$-MV$2$ reported by FCOS is $30.9\%$ with $90$K iterations. We double the training iterations and train with the distillation method. The final mAP on minival is $33.9\%$. The test results are in Table XIV, and we also list the AP and inference time for some state-of-the-art one-stage detectors to show the position of the baseline and our detectors trained with the structured knowledge distillation method. To make a fair comparison, we also double the training iterations without any distillation methods, and obtain mAP of $32.7\%$ on minival.

Conclusion

We have studied knowledge distillation for training compact dense prediction networks with the help of cumbersome/teacher networks. By considering the structure information in dense prediction, we have presented two structural distillation schemes: pair-wise distillation and holistic distillation. We demonstrate the effectiveness of our proposed distillation schemes on several recent compact networks on three dense prediction tasks: semantic segmentation, depth estimation and object detection. Our structured knowledge distillation methods are complimentary to traditional pixel-wise distillation methods.

Acknowledgements

C. Shen’s participation was in part supported by ARC DP Project “Deep learning that scales”.

References