Architecture Search of Dynamic Cells for Semantic Video Segmentation

TL;DR

This paper introduces a neural architecture search method to automatically design dynamic cells for semantic video segmentation, achieving stable, accurate results without optical flow computation in just 2 GPU-days.

Contribution

It proposes a novel NAS-based approach for dynamic cell architecture design in video segmentation, reducing reliance on manual design and optical flow.

Findings

Achieves high accuracy on CityScapes and CamVid datasets.

Requires only 2 GPU-days for architecture search.

Does not depend on optical flow computation.

Abstract

In semantic video segmentation the goal is to acquire consistent dense semantic labelling across image frames. To this end, recent approaches have been reliant on manually arranged operations applied on top of static semantic segmentation networks - with the most prominent building block being the optical flow able to provide information about scene dynamics. Related to that is the line of research concerned with speeding up static networks by approximating expensive parts of them with cheaper alternatives, while propagating information from previous frames. In this work we attempt to come up with generalisation of those methods, and instead of manually designing contextual blocks that connect per-frame outputs, we propose a neural architecture search solution, where the choice of operations together with their sequential arrangement are being predicted by a separate neural network. We showcase that such generalisation leads to stable and accurate results across common benchmarks, such as CityScapes and CamVid datasets. Importantly, the proposed methodology takes only 2 GPU-days, finds high-performing cells and does not rely on the expensive optical flow computation.

Tables

Table 1. Table 1: Description of operations used in the search process.

ID	Description
0	separable conv $3\times 3$
1	global average pooling followed by upsampling and conv $1\times 1$
2	separable conv $3\times 3$ with dilation rate $3$
3	separable conv $5\times 5$ with dilation rate $6$
4	skip-connection
5	deformable $3\times 3$ convolution

Table 2. Table 2: Description of aggregation operations used in the search process.

ID	Description
0	summation with per-channel learnable weights per each input
1	channel-wise concatenation of two inputs followed by conv $1\times 1$ to reduce the number of channels to the original size
2	(weight) predictive operation, where the first input becomes a set of spatial convolutional filters (weights) applied on the second one
3	bilinear sampling of the first input, where an affine grid is predicted based on the values of the second input [11]
4	3D-convolution where two inputs are stacked together forming a new dimension with $2\times 3\times 3$ convolution applied on top
5	dense attention: i.e. element-wise multiplication between the first input and the sigmoid-activated second one

Full text

1 Introduction

Human beings are well-equipped by evolution to quickly observe changes in dynamic environments. From merely few seconds of studying an unknown scene, we are able to coherently map out its main constituents. In contrast, static semantic segmentation networks would perform poorly in such conditions, and may as well produce contradictory predictions across the frames. Therefore, the question arises of how to make the static models suitable for segmenting continuously evolving scenes?

One well-known approach would be to use the optical flow that describes the motion in the scene between adjacent frames [9, 32]. The optical flow calculation tends to be expensive and also comes with several notable disadvantages, among which its inability to deal with occlusions and newly appeared objects. Nevertheless, as shown by Gadde et al. [9], a relatively poor estimate of the optical flow may still carry significant benefits, not the least of which lies in computational savings.

Alternatively, one may choose to model which information must be propagated across the frames, e.g. with the help of a recurrent neural network with memory units [20]. Even more biologically plausible are the models that compute different features at various time-scales [24], in a vein similar to neural spikes. Naturally, this comes with its own set of disadvantages, most notably the difficulty of choosing an appropriate scheduling regime for updating individual parts of the network.

Yet another complementary line of work focuses on approximating an expensive per-frame forward pass with cheaper alternatives: e.g. Li et al. [14] predicted local filters to be applied on the segmentation prediction from the previous frame, while Jain et al. [12] used a larger network for key frames and directly employed a smaller one for consecutive frames. Such savings may allow to re-use more expensive optical flow methods without a significant slowdown, but the choice of key frames can be crucial and not readily justifiable.

Looking closely at the aforementioned approaches for video semantic segmentation, one may notice an easily discernible pattern: a typical video segmentation network predicts a labelling of the current frame based on the information propagated from the previous one and hidden representations of the current one (Fig. 1). While seemingly obvious, it possesses certain variations depending on the goal - e.g. whether efficiency, or real-time performance is desired. Importantly, what we would like to emphasise here is that, while technically sound, all the current approaches have been manually designed and have not considered any interplay between different building blocks.

Starting from that general pattern we instead propose to leverage the neural architecture search (NAS) [33] methodology to find contextual blocks that enhance a per-frame segmentation network with dynamic components. This motivation is justified by recent results achieved using NAS on such tasks as image classification [34, 16], language modelling [21] and static semantic segmentation [5, 18], that oftentimes outperform manually designed networks. We build upon those results and adapt current approaches in a way suitable for handling the dynamic nature of dense per-pixel classification. To the best of our knowledge, we are the first to consider the application of NAS to the task of video semantic segmentation.

Our automated approach comes with certain benefits, concretely:

• i.)

it considers a larger span of initial building blocks than any previous work,

• ii.)

it empirically evaluates different design structures and finds most promising ones, and

• iii.)

it requires only few GPU-days to find a set of high-performing structures.

Furthermore, although we do not consider it in this work, the proposed methodology can further be extended to take into account different specific objectives (even non-differentiable), such as runtime [25].

2 Related Work

2.1 Static semantic segmentation

Most recent approaches in static semantic segmentation have been exploiting fully convolutional neural networks [17]. Typical methods are based either on the encoder-decoder structure with skip-connections [17, 15], dilated convolutional layers [28, 30, 6], or the combination of the above [7]. Per-frame instantiations of these networks are usually computationally expensive, hence, several works have considered building light-weight segmentation architectures [29, 19]. Nevertheless, due to the lack of information propagation between frames, these networks perform poorly on videos and are unable to provide consistent results.

2.2 Dynamic semantic segmentation

One of the first lines of work in video segmentation has been built upon the usage of the optical flow [32], in which features extracted from the previous frame are propagated to the current one via warping. This usually results in a slight computational overhead, although as noted by Gadde et al. [9] an easily attainable noisy estimate of the optical flow still carries significant benefits. Nevertheless, the optical flow does not fair well in situations when scenes are undergoing substantial changes with novel objects constantly appearing and multiple occlusions being present. Thus, Jain et al. [12] have proposed to combine the optical flow estimate with a relatively cheaper approximation of the current frame using a smaller network. Xu et al. [26] have chosen to assign different image regions to two different networks to process: while the first one - deep and slow - works on regions that have significantly changed, the second one - shallow - predicts new features based on the optical flow information. In a similar vein, Nilsson and Sminchisescu [20] have propagated labels from the previous frame at only those pixels where the optical flow estimate is reliable.

A seemingly different approach, proposed by Li et al. [14], instead predicts local convolutional kernels based on the low-level representation of the current frame that are applied on the prediction from the previous frame. Importantly, while the current estimate is being used for next frame, a more accurate one is being computed in parallel for future re-use.

In yet another line of work, Chandra et al. [3] have adapted Deep Gaussian Random Field [4] to handle temporal information by predicting besides unary and spatial pairwise terms also temporal pairwise terms, efficiently propagating features between frames.

2.3 Neural Architecture Search

NAS methods aim to find high-performing architectures in an automated way. Here, we consider the reinforcement learning-based (RL) approach [33], where a separate recurrent neural network (controller) outputs a sequence of tokens describing an architecture that should provide highest score on the holdout validation set.

While there is no prior work on NAS for video segmentation, two results in static segmentation are worth mentioning: Chen et al. [5] used a random search to find a single set of operations (so-called ‘cell’) on the top of the DeepLab architecture [6], while Nekrasov et al. [18] exploited RL to find a cell together with the topological structure of the encoder-decoder type of architecture. We borrow one of the architectures found by Nekrasov et al. as our static baseline, and extend their NAS approach for video segmentation. Since we are only searching for the dynamic component that connects different instantiations of the already pre-trained static segmentation network, we are able to train and evaluate each candidate in a short amount of time, the trait that is extremely important for all NAS methods.

3 Methodology

As noted in introduction and depicted in Fig. 1, we attempt to generalise previous solutions for video semantic segmentation in such a way that NAS methods become readily applicable. To this end, we look for a single cell that connects representations from the previous frame and enhances current predictions without a significant overhead. What follows is the description of the input space (Sect. 3.1), the search space (Sect. 3.2), and the search approach (Sect. 3.3).

3.1 Input space

We consider the arch2 network from the work of Nekrasov et al. [18]. It is an encoder-decoder type of the segmentation network with the encoder being a light-weight classifier (MobileNet-v2 [22]), and the decoder being an automatically discovered structure presented in Fig. 2. This architecture strikes a fine balance between accuracy and runtime, both being important characteristics for semantic video segmentation. Here it should be noted that the application of our methodology is not directly tied to a concrete architecture and can be easily adapted to work with other networks.

In the proposed setup, the static network is applied end-to-end on the first frame and three outputs are being recorded: an intermediate representation - in this case, the encoder’s output with the resolution of $\frac{1}{32}$ of the input image ($layer~{}4$), the decoder’s output before ($dec$) and after the final classifier ($pred$) - both with resolutions of $\frac{1}{8}$ of the input image and with $64$ and $C$ numbers of channels, correspondingly, where $C$ is the number of output classes. For the second frame, we record three outputs from the encoder only - two intermediate ones with the resolutions of $\frac{1}{8}$ ($layer~{}2$) and $\frac{1}{16}$ ($layer~{}3$), respectively, and the final one with the resolution of $\frac{1}{32}$ ($layer~{}4$).

We rely on the dynamic cell, the layout of which will be described below, to predict the semantic labelling of the current frame given $5$ inputs: $layer~{}4$ and $dec$ from the previous frame, and $layers~{}2{-}3{-}4$ from the current one. This way, we do not have to execute the decoder part of the static segmentation network on the current frame (thus decreasing latency), at the same time re-using information from the previous frame. The output of the dynamic cell serves as the input $dec$ for the next frame.

3.2 Search space

We rely on an LSTM-based controller to predict a sequence of operations together with locations where they should be applied in order to form a dynamic cell [18]. Concretely, we first choose two layers out of the provided five (with replacement), two corresponding operations that need to be applied on each of them, and an aggregation operation that combines two inputs into a single output. On the next step, we repeat this process, but now we are sampling two layers out of six possible, with the aggregated result being added into the sampling pool. This process can be repeated multiple times, with the final output being formed by the concatenation of all non-sampled aggregated results.

We rely on a similar set of operations as for static segmentation (Table 1), and in order to enable the dynamic cell to apply convolutional filters on irregular grids, we also include deformable $3\times 3$ convolution [31].

While Nekrasov et al. [18] simply summed up two different inputs at each step, here to compensate for the dynamic nature of our problem we consider a set of aggregation operations given in Table 2.

Based on the previous works, we conjecture that this set of operations will be sufficient for the task of video segmentation, and we provide experimental results to support this claim.

3.3 Finding optimal architectures

We assume that there exists a video dataset that comes with segmentation annotations for at least a subset of consecutive frames. From it, we build pairs (or triplets) of frames such that in each sequence all the frames following the first one are always annotated. As commonly done, we further divide this set into two disjoint parts - meta-train and meta-val. We further assume an existence of the static segmentation network pre-trained on this dataset111Please refer to Sect. 6 for the details on pre-training of static segmentation networks. - in particular, arch2 from [18]. As mentioned above, we chose this particular architecture due to its compactness and low latency.

The controller samples a structure of the dynamic cell which we train on the meta-train set and evaluate on meta-val. As done in [18], we consider the geometric mean of three metrics as the validation score: mean intersection-over-union (mIoU), frequency-weighted IoU (fwIoU) and mean-pixel accuracy (mAcc). This score is used by the controller to update its weights, and the process is repeated multiple times. After that, one can either sample several cells from the trained controller, or simply choose best found cells that achieved highest results during the search process.

4 Experiments

We conduct all our experiments on two popular video segmentation benchmark datasets - CamVid [2] and CityScapes [8].

The first one, CamVid, comprises $701$ RGB images of resolution $480{\times}360$ densely annotated into $11$ categories. Following previous work [1], we use the dataset splits of $367$ images for training, $101$ - for validation and $233$ - for testing. We train generated architectures with batches of examples each comprising $3$ consecutive frames.

The CityScapes dataset contains $5000$ high-resolution $2048{\times}1024$ images densely labelled with $19$ semantic classes - $2975$ for training, $500$ for validation and $1525$ for testing, respectively. In addition to that, raw unannotated frames extracted from videos are also provided. For each annotated example, we add an image frame that precedes it and train architectures with batches of sequences of length $2$, in which the second frame is always annotated.

In each case, we initialise the decoder’s output dec on the first frame using the pre-trained static segmentation network, and rely on the dynamic cell at all following frames in the sequence as described in Sect. 3.1. To update the dynamic cell weights, we sum up cross-entropy loss terms at each frame after the first one and back-propagate the gradients.

For both, search and training, we exploit a single V$100$ GPU with $32$GB of memory.

4.1 Search

For searching we only employ the training splits of each dataset. We further divide each randomly in $2$ non-overlapping sets - meta-train ($90\%$) and meta-val ($10\%$). We pre-compute all required outputs from the pre-trained static network and store them in memory. The static network is kept unchanged during the whole search process. Each generated architecture is trained on the meta-train split and evaluated on meta-val. We keep track of average performance and apply early stopping halfway through the training if the generated architecture is un-promising as done in [18].

Our controller is a two-layer LSTM with $100$ hidden units randomly initialised from uniform distribution [18]. The controller is trained with PPO [23] with the learning rate of $1{e-}4$. To reduce the size of generated cells, we set the number of emitted layers (each layer is a string of five tokens as described in Sect. 3.2) to $5$ on CamVid and to $4$ on CityScapes.

For CamVid, we train predicted cells on mini-batches of $48$ sequences for $20$ epochs with the learning rate of $8{e-}3$ and the Adam learning rule [13]. Each image–segmentation mask pair in the sequence is cropped to $600$ with the shorter side being mean-padded to $400$. No transformations are applied to the validation sequences.

For CityScapes, we train for $10$ epochs with $48$ sequences each cropped to $512{\times}512$ with the longer side being resized to $1024$.

Results

We visualise the progress of rewards on each dataset in Fig. 3. Although the rewards are not directly comparable between the datasets, the growth dynamics on both datasets signal that the controller is able to discover better architectures throughout the search process.

We further look at the distributions of sampled operations, aggregation operations and input layers plotted on Fig. 4. On both datasets, global average pooling and separable $5{\times}5$ convolution with dilation rate $6$ are sampled less frequently than other operations, potentially indicating that these layers could be omitted from the search process. On average, the controller trained on CityScapes prefers sampling deformable convolution (Fig. 4a), while the CamVid one - separable $3{\times}3$ convolution (Fig. 4d).

In terms of aggregation operations, the dynamics between two controllers vary significantly: the CamVid-based controller tend to rely on dense attention, while omitting the predictive operation (Fig. 4e). In contrast, the CityScapes controller is more likely to apply bilinear sampling on an affine grid, and to ignore predictive operation together with dense attention (Fig. 4b).

When sampling the input layers, the controllers behave similarly: in particular, both tend to skip $layer~{}{4}$ from the previous and current frames. The CityScapes controller extensively uses information from the previous $dec$ layer (Fig. 4c), while the CamVid one - from $layer~{}{2}$ of the current frame (Fig. 4f). This may well imply that on CityScapes the final predictions on the current frame change only slightly with respect to the previous frame.

Importantly, these observations indicate that two controllers trained on two different datasets exhibit various patterns, potentially capturing dataset-specific attributes in order to discover better performing architectures.

4.2 End-to-end Training

We further select top-$2$ performing dynamic cells on each dataset to train end-to-end on full training sets for longer.

In particular, for CamVid, we pre-train the dynamic cell only with Adam and the learning rate of $8{e-}3$ for $10$ epochs with the batch size of $16$ sequences. Then we decrease the cell’s learning rate in half, and fine-tune the whole architecture (i.e., with the per-frame segmentation network) end-to-end for $100$ epochs - the static network weights are updated using SGD with momentum of $0.9$ and the learning rate of $5{e-}4$. Each sample in the batch is cropped to $600{\times}600$ with the shorter side being padded to $400$.

On CityScapes we pre-train for $200$ epochs with the batch size of $16$ sequences and fine-tune end-to-end for $200$ epochs. Each example in the batch is cropped to $769{\times}769$.

CamVid Results

We provide quantitative results on CamVid in Table 3. The inclusion of dynamic cells in both cases leads to an improvement over baseline by more than $1\%$. Importantly, with the exclusion of first frame in the sequence, we do not rely on expensive computations involving the static decoder.

Both our models perform comparably to other state-of-the-art video segmentation networks even though the backbone that we rely on - MobileNet-v2 [22] - is much smaller in comparison to ResNet-101 [10] exploited by Chandra et al. [3], or DilatedNet [27] - by Gadde et al. [9] and GRFP [20]. Furthermore, we did not make any use of higher-resolution images of $960{\times}720$ to further improve our scores.

We further visualise a few qualitative examples in Fig. 5. The dynamic cell enables the network to effectively propagate information about thin structures, such as poles, which makes the resultant segmentation masks consistent in contrast to the per-frame baseline (rows $1{-}5$). Furthermore, the multi-frame segmentation network is able to track objects across neighbouring frames (rows $1{-}2$).

CityScapes Results

We include the validation results of two discovered cells on CityScapes in Table 4. Once again, both dynamic cells are able to outperform the per-frame baseline by $1.2\%$. Furthermore, our models achieve favourable results in comparison to other video segmentation methods, all of which employ significantly larger backbones and, with the exclusion of Li et al. [14], all rely on the optical flow computation. Note also that Gadde et al. [9] improved over their respective static baseline by $1.2\%$, too, while introducing a non-negligible overhead of $40$ms; and Li et al. [14] compromised more than $3\%$ of the baseline score in order to reduce the latency. In contrast, we overcame our static baseline and decreased the runtime (Table 5).

A few inference examples are visualised in Fig. 6. As can be seen, the dynamic cells enhance the per-frame baseline results and identify partially occluded vehicles more accurately (rows $1{-}2$, $5$), while also avoiding misclassification of traffic signs at pixels with similar texture patterns (rows $2{-}4$).

4.3 Details of Discovered Architectures

We include characteristics of our networks together with numbers reported by others in Table 5. As evident, our dynamic segmentation approach is superior to others in terms of its latency and compactness. Concretely, all our architectures contain at most $3.4$M parameters while having an average per-frame runtime of $50$ms on high-resolution $2048{\times}1024$ images. This is possible due to both the network design and the exclusion of the optical flow computation.

All the trained cells are visualised in Fig. 7. Notably, layers with deformable convolution are present in all architectures. To propagate information from the previous frame, each cell exploits the $dec$ output instead of $layer~{}4$. All the cells prefer aggregating outputs via channel-wise concatenation with cell0 also relying on dense attention, and cell3 – on affine transformation with bilinear sampling. In addition, cell1 and cell2 employ $3$D convolution in order to capture information between various inputs.

5 Discussion & Conclusions

It is still an open question of what is the optimal way of propagating and extracting information across video frames. While a straightforward solution involving the optical flow allows to achieve solid results, it possesses several disadvantages that stem from the limitations of the optical flow itself and ultimately limit the ability of the network to adapt to novel frames. Furthermore, computations involving the optical flow cause a significant overhead, prohibiting the final system from being deployed in real-time.

In this work, instead of manually enhancing static segmentation networks with dynamic components, we proposed an automatic approach based on neural architecture search methods. Such automation have multiple benefits as it explores a large pool of networks and finds best-performing ones on the given dataset. In a broader sense, starting from a static per-frame segmentation network, we showcased a way of generalising existing solutions without any reliance on the optical flow. In particular, we extended the static baseline with a dynamic cell, the design of which is automatically discovered with the help of reinforcement learning. The best discovered cells improve the baseline by more than $1\%$ at the same time leading to significant memory and latency savings. Concretely, two discovered cells on CityScapes reach $75.6\%$ mean IoU and require only $50$ms on average to process a $2048{\times}1024$ frame. While the proposed methodology relies on the static baseline, we expect that omitting that requirement and searching for a video segmentation network end-to-end would be an interesting problem to consider in the future work.

Acknowledgements

VN, CS, IR’s participation in this work were in part supported by ARC Centre of Excellence for Robotic Vision. CS was also supported by the GeoVision CRC Project.

Appendix

6 Training Details of Static Baseline

The static baseline that we consider in the main text is arch2 from [18], which we pre-train on CamVid [2] and CityScapes [8].

We utilise the ‘poly’ learning schedule [6] with the initial learning rates of $5{e-}2$ and $1{e-}2$ for the encoder and the decoder, respectively. As in [18], we set the weight for auxiliary losses to $0.3$.

On CityScapes, we train for $1000$ epochs with mini-batches of $28$ examples each randomly scaled with the scale factor in range of $[0.5,2.0]$ and randomly cropped to $800{\times}800$ with each side zero-padded accordingly. On CamVid, we train for $2000$ epochs with mini-batches of $32$ examples each randomly scaled with the scale factor in range of $[0.5,2.0]$ and randomly cropped to $600{\times}600$ with each side zero-padded accordingly.