R EINFORCED A CTIVE L EARNING FOR I MAGE S EGMENTATION
A BSTRACT
Learning-based approaches for semantic segmentation have two inherent chal
- lenges.
First, acquiring pixel
-wise labels is expensive and time
-consuming
Second, realistic segmentation datasets are highly unbalanced: some categories are much more abundant than others biasing the performance to the most represented ones.
In this paper, we are interested in focusing human labelling effort on a small subset of a larger pool of data, minimizing this effort while maximizing performance of a segmentation model on a hold-out set.
We present a new active learning strategy for semantic segmentation based on deep reinforcement learning (RL).
An agent learns a policy to select a subset of small informative image regions – opposed to entire images – to be labeled, from a pool of unlabeled data.
The region selection decision is made based on predictions and uncertainties of the segmentation model being trained.
Our method proposes a new modiﬁcation of the deep Q-network (DQN) formulation for active learning, adapting it to the large-scale nature of semantic segmentation problems.
We test the proof of concept in CamVid and provide results in the large-scale dataset Cityscapes.
On Cityscapes, our deep RL region-based DQN approach requires roughly 30% less additional labeled data than our most competitive baseline to reach the same performance.
Moreover, we ﬁnd that our method asks for more labels of under-represented categories compared to the baselines, improving their performance and helping to mitigate class imbalance
I NTRODUCTION
Semantic segmentation, the task of labelling an image pixel-by-pixel with the category it belongs to, is critical for a variety of applications such as autonomous driving (Müller et al., 2018; Wang & Pan, 2018), robot manipulation (Schwarz et al., 2018), embodied question answering (Yu et al., 2019) and biomedical image analysis (Ronneberger et al., 2015).
Convolutional neural networks (Lecun et al.
1998)-based methods have achieved excellent results on large-scale supervised semantic segmentation, in which we assume pixel-level annotations are available (Farabet et al., 2013; Pinheiro & Collobert, 2014; Long et al., 2015).
For such models to work, however, they need a large amount of pixel-level annotations that may require costly human labor (Cordts et al., 2016; Bearman et al., 2016).
Current semantic segmentation datasets have pixel-wise annotations for each image.
This standard approach has two important issues: (i) pixel-level labelling is extremely time consuming.
For example, annotation and quality control required more than 1.5h per image (on average) on Cityscapes (Cordts et al., 2016), a popular dataset used for benchmarking semantic segmentation methods.
(ii) Class imbalance in the data is typically extreme.
Certain categories (such as ‘building’ or ‘sky’) can appear with two orders of magnitude more frequently than others ( e.g. ‘pedestrian’ or ‘bicycle’).
This can lead to undesired biases and performance properties for learned models.
This is specially relevant when we want to collect annotated data with a human in the loop to create a new dataset or to add more labeled data to an existing one.
We can tackle the aforementioned problems by selecting, in an efﬁcient and effective way, which regions of the images should be labeled next.
Active learning (AL) is a well-established ﬁeld that studies precisely this: selecting the most informative samples to label so that a learning algorithm will perform better with less data than a non-selective approach, such as labelling the entire collection of data.
Active learning methods can be roughly divided in two groups: (i) methods that combine different manually-designed
AL strategies (Roy & McCallum, 2001; Osugi et al., 2005; Gal et al., 2017; Baram et al., 2004;
Chu & Lin, 2016; Hsu & Lin, 2015; Ebert et al., 2012; Long & Hua, 2015) and (ii) data-driven
AL approaches (Bachman et al., 2017; Fang et al., 2017; Konyushkova et al., 2017; Woodward &
Finn, 2016; Ravi & Larochelle, 2018; Konyushkova et al., 2018), that learn which samples are most informative to train a model using information of the model itself.
Although label acquisition for semantic segmentation is more costly and time consuming than image classiﬁcation, there has been considerably less work in active learning for semantic segmentation (Dutt Jain & Grauman, 2016;
Mackowiak et al., 2018; Vezhnevets et al., 2012; Konyushkova et al., 2015; Gorriz et al., 2017;
Yang et al., 2017), and they focus on hand-crafted strategies.
Current AL techniques that use reinforcement learning (Konyushkova et al., 2018; Fang et al., 2017;
Woodward & Finn, 2016; Pang et al., 2018; Padmakumar et al., 2018; Bachman et al., 2017) focus on labelling one sample per step until a budget of labels is met.
In semantic segmentation, this would translate into labelling a single region per step.
This is highly inefﬁcient, since each step involves updating the segmentation network and computing the rewards.
In this work, we propose an end-to-end method to learn an active learning strategy for semantic segmentation with reinforcement learning by directly maximizing the performance metric we care about, Intersection over Union (IoU).
We aim at learning a policy from the data that ﬁnds the most informative regions on a set of unlabeled images and asks for its labels, such that a segmentation network can achieve high-quality performance with a minimum number of labeled pixels.
Selecting regions, instead of entire images, allows the algorithm to focus on the most relevant parts of the images, as shown in Figure 1.
Although class imbalance in segmentation datasets has been previously addressed in (Badrinarayanan et al., 2017; Chan et al., 2019; Sudre et al., 2017), among others, they try to solve a problem that arises from the data collection process.
We show that our proposed method can help mitigate the problem at its source, i.e. in the data annotation itself.
Because our method maximizes the mean IoU per class, it indirectly learns to ask for more labels of regions with under-represented classes, compared to the baselines.
Moreover, we propose and explore a batch-mode active learning approach that uses an adapted DQN to efﬁciently chose batches of regions for labelling at each step.
To the best of our knowledge, all current approaches for active learning in semantic segmentation rely on hand-crafted active learning heuristics.
However, learning a labelling policy from the data could allow the query agent to ask for labeled data as a function of the data characteristics and class imbalances, that may vary between datasets.
Our main contributions can be summarized as follows:
(i
) we learn a RL-based acquisition function for region-based active learning for segmentation, (ii) we formulate our active learning framework with a batch-mode DQN, which labels multiple regions in parallel at each active learning iteration (a more efﬁcient strategy for large-scale datasets that is compatible with standard mini-batch gradient descent), and (iii) we test the proof of concept in dataset and provide results
2016) dataset, beating a recent state-of-the-art technique known as BALD (Gal et al., 2017), a widely used entropy-based selection criterion and uniform sampling baselines.
R ELATED W ORK
Active learning .
Traditional active learning techniques focus on estimating the sample informativeness using hand-crafted heuristics derived from sample uncertainty:
employing entropy (Shannon, 1948), query-by-committee (Dagan & Engelson, 1995; Shannon, 1948; Freund et al., 1993), maximizing the error reduction (Roy & McCallum, 2001), disagreement between experts (Dagan &
Engelson, 1995; Freund et al., 1993) or Bayesian methods that need to estimate the posterior distribution (Houlsby et al., 2011a; Gal et al., 2017).
Some approaches combine different techniques to improve AL performance.
For instance, relying on exploration-exploitation trade-offs (Osugi et al., 2005), on a bandit formulation (Baram et al., 2004; Chu & Lin, 2016; Hsu & Lin, 2015) and on reinforcement learning (Ebert et al., 2012; Long & Hua, 2015).
However, these methods are still limited in the sense that they combine hand-crafted strategies instead of learning new ones.
More recent active learning methods rely on an acquisition function that estimates the sample informativeness with a learned metric.
Konyushkova et al.
(2017) estimate the error reduction of labelling a particular sample, choosing the ones that maximize the error reduction.
Wang et al.
(2017) introduce a cost-effective approach that also uses conﬁdent predictions as pseudo ground truth labels.
AL with reinforcement learning.
Recently, reinforcement learning has gained attention as a method to learn a labelling policy that directly maximizes the learning algorithm performance.
For instance, Liu et al.
; Bachman et al.
(2017) leverage expert knowledge from oracle policies to learn a labelling policy, and Pang et al.
; Padmakumar et al.
(2018) rely on policy gradient methods to learn the acquisition function.
In a different approach, some methods gather all labeled data in one big step.
In Contardo et al.
(2017), all samples are chosen in one step with a bi-directional
RNN for the task of one-shot learning.
In Sener & Savarese (2018), they propose to select a batch of representative samples that maximize the coverage of the entire unlabeled set.
However, the bounded core-set loss used tends to perform worse when the number of classes grows.
More similar to our approach, some prior works propose to learn the acquisition function with a Deep
Q-Network (DQN) (Mnih et al., 2013) formulation.
These works have examined both stream-based active learning (Fang et al., 2017; Woodward & Finn, 2016), where unlabeled samples are provided one by one, and the decision is to label it or not, and pool-based active learning (Konyushkova et al., 2018), where all the unlabeled data is provided beforehand, and the decision is later taken on which samples to choose.
The work of Konyushkova et al.
(2018) is the closest to ours.
Similar to them, our method also leverages the beneﬁts of Q-learning (Watkins & Dayan, 1992) to tackle pool-based AL.
Contrary to them, we deal with a much more complex problem: semantic segmentation versus simple classiﬁcation on UCI repository (Dua & Graff, 2017).
The large-scale nature of the problem requires us to use a very different deﬁnition of actions, states and rewards.
Moreover, we need to adapt the
DQN formulation to allow the problem to be computationally feasible.
AL for semantic segmentation.
Active learning for semantic segmentation has been relatively less explored than other tasks, potentially due to its large-scale nature.
For instance, Dutt Jain &
Grauman (2016) combine metrics (deﬁned on hand-crafted heuristics) that encourage the diversity and representativeness of labeled samples.
Some rely on unsupervised superpixel-based oversegmentation (Vezhnevets et al., 2012; Konyushkova et al., 2015) – and highly depend on the quality of the super-pixel segmentation.
Others focus on foreground-background segmentation of biomedical images (Gorriz et al., 2017; Yang et al., 2017), also using hand-crafted heuristics.
Settles et al.
Vijayanarasimhan & Grauman (2009); Mackowiak et al.
(2018) focus on cost-effective approaches, proposing manually-designed acquisition functions based on the cost of labeling images or regions of images.
However, this information is not always given, restricting their applicability.
Mackowiak et al.
(2018) focus on cost-effective approaches, where the cost of labeling an image is not considered equal for all images.
Similar to our work, they use a region-based approach to cope with the large number of samples on a segmentation dataset.
Contrary to us, their labelling strategy is based on manually deﬁned heuristics, limiting the representability of the acquisition function.
To the best of our knowledge, our work is the ﬁrst to apply data-driven RL-based approach to the problem of active learning for semantic segmentation.
M ETHOD
We are interested in selecting a small number of regions 1 (cropped from images in the original dataset) from a large unlabeled set to maximize the performance of a segmentation network f , parameterized by θ .
This process is done iteratively until a given budget B of labeled samples is achieved.
At each iteration t , a query network π , parameterized by φ , selects K regions to be labeled by an oracle from a large unlabeled set U t .
These samples are added to the labeled set L t , that is used to train the segmentation network f .
The performance is measured with a standard semantic segmentation metric, Intersection-over-Union (IoU).
We cast the AL problem within a Markov decision process (MDP) formulation, inspired by other work such as (Padmakumar et al., 2018;
Fang et al., 2017; Bachman et al., 2017; Pang et al., 2018;
Konyushkova et al., 2018).
We model the query network π as a reinforcement learning agent, speciﬁcally a deep Q-network (Mnih et al., 2013).
This data-driven approach allows the model to learn selection strategies based solely on prior AL experience.
Our formulation differs from other approaches by the task we address, the deﬁnitions of states, actions and rewards, and the reinforcement learning algorithm
we use to ﬁnd the optimal policy.
3. A CTIVE L EARNING WITH R EINFORCEMENT L EARNING FOR S EGMENTATION
In our setting, we use four different data splits.
To train π , we deﬁne a subset of labeled data D T to play the active learning game for several episodes and learn a good acquisition function that maximizes performance with a budget of B regions.
The query network is evaluated on a different split D V .
We use a separate subset D R to obtain the reward signal by evaluating the segmentation network on it.
The set D S ( |D S | (cid:28) |D T | ) is used to construct the state representation.
The MDP is deﬁned with the sequence of transitions { ( s t , a t , r t +1 , s t +1 ) } .
For every state s t ∈ S (function of the segmentation network at timestep t ), the agent can perform actions a t ∈ A to choose which samples from U t to annotate.
The action a t = { a kt } Kk =1 , composed of K sub-actions, is a function of the segmentation network, the labeled and the unlabeled set.
Each sub-action asks for a speciﬁc region to be labeled.
Then, it receives a reward r t +1 based on the improvement in mean
IoU per class after training the segmentation network with the selected samples.
Note that states and actions do not depend on the speciﬁc architecture of the segmentation network.
We are interested in ﬁnding a policy to select samples that maximize the segmentation performance.
We use deep
Q-network (Mnih et al., 2013) and samples from an experience buffer E to train the query network π .
Each episode e elapses a total of T steps.
We start by setting the segmentation network f to a set of initial weights θ 0 and with no annotated data, i.e. , L 0 = ∅
and U
= D T .
At each iteration t , the following steps are done:
The state s t is computed as function of f t and D S .
A restricted action space is built with K pools P kt with N regions, sampled uniformly from the unlabeled set U t .
For each region in each pool, we compute its sub-action representation a k,nt .
The query agent selects K sub-actions { a kt } Kk =1 with (cid:15) -greedy policy.
Each sub-action a kt is deﬁned as selecting one region x k (out of N ) to annotate from a pool P kt .
An oracle labels the regions and the sets are updated: L t +1 =
L t ∪ { ( x k , y k ) } Kk =1 and U t +1 =
U t \ { x k } Kk =1 . 5.
The segmentation network f t +1 is trained one iteration on the recently added regions { x k } Kk =1 . 6.
The agent receives the reward r t +1 as the difference of performance between f t +1 and f t on D R
Figure 2 depicts this training algorithm.
We consider the termination of each episode when the budget
B of labeled regions is met, i.e. , |L t | =
B .
Once the episode is terminated, we restart the weights of the segmentation network f to the initial weights θ 0 , set L 0 = ∅
and U
= D T , and restart the episode.
We train the query policy π by simulating several episodes and updating its weights at each
[Figure]
timestep by sampling transitions { ( s t , a t , r t +1 , s t +1 ) } from the experience replay buffer E .
More details in Section 3.2.
State representation.
We would like to use the state of the segmentation network f as the MDP state.
Unfortunately, it is not straightforward to embed f into a state representation.
Following Konyushkova et al.
(2017), we represent the state space S with the help of a set-aside set D S .
We use a small subset of data from the train set, making sure it contains a signiﬁcant representation of all classes.
We consider it to be a representative set of the dataset, and that any improvement in the segmentation performance on subset D S will translate into an improvement over the full dataset 2 .
We use the predictions of the segmentation network f t on D S to create a global representation of state s t (step 1 in Figure 2).
We need a compact representation to avoid intensive memory usage due to the pixel-wise predictions.
The samples in D S are split in patches, and compact feature vectors are computed for all of them.
Then, each region is encoded by the concatenation of two sets of features: one is based on class predictions of f t and the other on its prediction uncertainty, represented as the Shannon entropy (Shannon, 1948).
The ﬁrst set of features (i) is a (normalized) count of the number of pixels that are predicted to each category.
This feature encodes the segmentation prediction on a given patch while dismissing the spatial information, less important for small patches.
Moreover, we measure the uncertainty of the predictor with the entropy over the probability of predicted classes.
For each region, we compute the entropy of each pixel location to obtain a spatial entropy map.
To compress this representation, we apply min, average and max-poolings to the entropy map to obtain downsampled feature maps.
The second set of features (ii) is thus obtained by ﬂattening these entropy features and concatenating them.
Finally, the state s t is represented by an ensemble of the feature representation of each region in D S .
Figure A.1a illustrates how s t is computed from each region.
Action representation.
In our setting, taking an action means asking for the pixel-wise annotation of an unlabeled region.
Due to the large-scale nature of semantic segmentation, it would be prohibitively expensive to compute features for each region in the unlabeled set at each step.
For this reason, instead, at each step t , we approximate the whole unlabeled set by sampling K pools of unlabeled regions P kt , each containing N (uniformly) sampled regions.
For each region, we compute its sub-action representation a k,nt (step 2 in Figure 2).
Each sub-action a k,nt is a concatenation of four different features: the entropy and class distribution features (as in the state representation), a measure of similarity between the region x k and the labeled set and another between the region and the unlabeled set.
The intuition is that the query network could learn to build a more class-balanced labeled set while still taking representative samples from the unlabeled set.
This could help mitigate the hard imbalance of the segmentation datasets and improve overall performance.
For each candidate region, x in a pool P kt , we compute the KL divergence between the class distributions of the prediction map of region x (estimated as normalized counts of predicted pixels in each category) and the class distributions of each labeled and unlabeled regions (using the groundtruth annotations and network predictions, respectively).
For the labeled set, we compute a KL divergence score between each of the labeled regions’ class distribution and the one of region x .
Summarizing all these KL divergences could be done by taking the maximum or summing them.
However, to obtain more informative features, we compute a normalized histogram of KL divergence scores, resulting in a distribution of similarities.
As an example, if we were to sum all the scores, having half of the labeled regions with a KL divergence of zero and the other half with a value c , would be equivalent to have all labeled regions with a KL divergence of c/ 2 .
The latter could be more interesting, since it means there are no labeled regions with the same class distribution as x .
For the unlabeled set we follow the same procedure, resulting in another distribution of KL divergences.
Both of them are concatenated and added to the action representation.
Figure A.1b illustrates how we represent each possible action in a pool.
Based on early experimentation, learning the state and action representations directly with a CNN does not provide strong enough features for the reinforcement learning framework to converge to a good solution.
An ablation study on the state and action components can be found in Appendix E.1.
3. B ATCH MODE DQN
The desired query agent should follow an optimal policy.
This policy maps each state to an action that maximizes the expected sum of future rewards.
We rely on a DQN (Mnih et al., 2013), parameterized by φ , to ﬁnd an optimal policy.
We train our DQN with a labeled set D T and compute the rewards in a held-out split D R .
As mentioned above, the query agent in our method selects K regions before transitioning to the next state.
We assume that each region is selected independently, as in the case where K annotators label one region in parallel.
In this case, the action a t is composed of K independent sub-actions { a kt } Kk =1 , each with a restricted action space, avoiding the combinatorial explosion of the action space.
To ease computation and avoid selecting repeated regions in the same time-step, we restrict each sub-action a kt to select a region x k in P kt deﬁned as:
[Equation]
for each k ∈ { 1 , ..., K } action take in timestep t .
The network is trained by optimizing a loss based on temporal difference (TD) error (Sutton, 1988).
The loss is deﬁned as the expectation over decomposed transitions T k
= { ( s t , a kt , r kt +1 , s t +1 ) } , obtained from the standard transitions { ( s t , a t , r t +1 , s t +1 ) } , by approximating
[Equation]
where E is the experience replay buffer and y kt the TD target for each sub-action.
To stabilize the training, we used a target network with weights φ (cid:48) and the double DQN (Van Hasselt et al., 2016) formulation.
The action selection and evaluation is decoupled; the action is selected with the target network and is evaluated with the query network.
The TD target for each sub-action is represented as:
[Equation]
where γ is a discount factor.
This formulation is valid under the approximation that the sub-actions are independent of each other, conditioned on the state.
We observed that increasing the number of sub-actions K per step eases computation and does not hinder segmentation performance.
We provide an ablation study on the effect of K in Appendix E.3.
E XPERIMENTS
We start this section by describing the datasets that we use to evaluate our method, the experimental setup, and the baselines.
We evaluate the algorithm in Camvid as a proof of concept and we show large-scale results on Cityscapes.
4. E XPERIMENTAL S ETUP
Although we can apply active learning in a setting with unlabeled data with a human in the loop that labels selected regions, we test our approach in fully labeled datasets, where it is easier to mask out the labels of a part of the data and reveal them when the active learning algorithm selects them.
CamVid (Brostow et al., 2008).
This dataset consists of street scene view images, with the resolution of 360 × 480 and 11 categories.
It has 370, 104 and 234 images for train, validation and test set, respectively.
We split the train set with uniform sampling in 110 labeled images (from where we get 10 images to represent the state set D S and the rest for D T ), and 260 images to build D V , where we evaluate and compare our acquisition function to the baselines.
The state set is chosen to be representative of D T , by restricting the sampling of D S to have a similar class distribution to the one of D T .
Each image is split into 24 regions of dimension 80 × 90.
We use the dataset’s validation set for D R .
We report the ﬁnal segmentation results on the test set.
In our experiments, we chose
K = 24 regions per step.
Our model is quite robust to the number of regions selected at each time step (see Appendix E.3).
Cityscapes (Cordts et al., 2016).
It is also composed of real street scene views, with image resolution of 2048 × 1024 and 19 semantic categories.
The train set with ﬁne-grained segmentation labels has 2975 images and the validation dataset of 500 images.
We uniformly sampled 360 labeled images from the train set.
Out of these, 10 images represent D S , 150 build D T and 200, D R , where we get our rewards.
The remaining 2615 images of train set are used for D V , as if they were unlabeled.
We report the results in the validation set (test set not available).
Each image is split in 128 regions of dimension 128 × 128.
We chose K = 256 regions per step.
Implementation details.
The split D R is used to get the rewards for the DQN and also for hyperparameter selection, that are chosen according to the best setup for both baselines and our method.
We report the average and standard deviation of the 5 different runs (5 random seeds).
As data augmentation, we use random horizontal ﬂips and random crops of 224 × 224 .
For more details, please refer to Appendix B on supplementary material.
Evaluation.
The query network π is trained on D T with a small, ﬁxed budget (0.5k regions for
Camvid and 4k regions for Cityscapes) to encourage picking regions that will boost the performance in an heavily scarce data regime.
The learned acquisition function, as well as the baselines, is evaluated on D V , where we ask for labels until the budget is met, for different budgets.
Note that the baselines do not have any learnable component.
Once the budget is reached, we train the segmentation network f with L T until convergence (with early stopping in D R ).
For a fair comparison, all methods’ segmentation network has been pre-trained (initial f weights θ 0 ) on GTA dataset (Richter et al., 2016), a synthetic dataset where high amounts of labeled data can be obtained without human effort, and D T (where we had labels to train the DQN).
We evaluate the ﬁnal segmentation performance (measured in mean IoU) on the test set of CamVid and on the validation set of Cityscapes.
4. R ESULTS
[Figure]
Results in CamVid.
We compare our results against three distinct baselines: (i) U is the uniform random sampling of the regions to label at each step out of all possible regions in the pool, (ii) H is an uncertainty sampling method that selects the regions with maximum cumulative pixel-wise Shannon entropy, (iii) B picks regions with maximum cumulative pixel-wise BALD (Houlsby et al., 2011b; Gal et al., 2017) metric.
We use 20 iterations of MC-Dropout (Gal & Ghahramani, 2016)
(instead of 100, as in (Gal et al., 2017)) for computational reasons.
In preliminary experiments, we did not observe any improvement using over 20 iterations.
In Camvid, we use a pool size of 10 for our method, H , B and 50 for U .
In Cityscapes, we have access to more data so we use pool sizes of 500, 200, 200 and 100 respectively for U , H ,
B and our method.
Pool sizes were selected according to the best validation mean
IoU.
Figure 4a shows results on CamVid for different budget sizes.
Our method outperforms the baselines for every ﬁxed budget, except for 1.5k regions, where we achieve similar performance as H .
We argue that the dataset has a small number of images and selecting 1.5k regions already reaches past 98% of maximum performance, where differences between our method and H are negligible.
Surprisingly, B is worse than U , specially for small budgets, where training with the newly acquired labels does not provide any additional information.
It overﬁts quickly to the training, getting a worst result that with the initial weights.
In general, all results have a high variance due to the low regime of data we are working in.
In Appendix E.2 we show the advantages of labeling small regions instead of full images.
Results in Cityscapes.
Figure 4b shows results on Cityscapes for different budgets.
Here, we also observe that our method outperforms the baselines for all budgets points.
Labelling 20k regions, corresponding to only 6% of the total pixels (additional to the labeled data in D T ), we obtain a performance of 64.5% mean IoU. This is 96% of the performance of the segmentation network if it had access to all labeled pixels.
To reach the same performance, H requires an additional 6k labeled regions (around 30% more pixels, equivalent to an extra 45 images).
In this larger dataset,
B performs better than random, showing that for the task of segmentation, B might start to show its beneﬁts only for considerably large budgets.
Table 1 shows the per-class IoU for the evaluated methods (with a ﬁxed budget).
Our method works specially well for under-represented classes, such as Person, Motorcycle or Bicycle, among others.
Indeed, our method selects more pixels belonging to under-represented classes than baselines.
Note that this is a side effect of directly optimizing for the mean IoU and deﬁning class-aware representations for states and actions.
Figure 3 shows the entropy of the distribution of selected pixels of the ﬁnal labeled set (for a budget of 12k regions) for
Cityscapes.
The higher the entropy means closer to uniform distribution over classes, and our method
[Table]
has the highest entropy.
Appendix C shows the distribution from which the entropy is computed and
Appendix D presents some qualitative results, showing what each method decides to label for some images.
C ONCLUSION
We propose a data-driven, region-based method for active learning for semantic segmentation, based on reinforcement learning.
The goal is to alleviate the costly process of obtaining pixel-wise labels with a human in the loop.
We propose a new modiﬁcation of DQN formulation to learn the acquisition function, adapted to the large-scale nature of semantic segmentation.
This provides a computationally efﬁcient solution that uses less labeled data than competitive baselines, while achieving the same performance.
Moreover, by directly optimizing for the per-class mean IoU and deﬁning class-aware representations for states and actions, our method asks for more labels of under-represented classes compared to baselines.
This improves the performance and helps to mitigate class imbalance.
As future work, we highlight the possibility of designing a better region deﬁnition, that could help improve the overall results, and adding domain adaptation for the learnt policy, to transfer it between datasets. A CKNOWLEDGEMENTS
We thank NSERC and PROMPT.
We would also like to thank the team at ElementAI for supporting this research and providing useful feedback.
A S TATE AND ACTION REPRESENTATION DETAILS
In this section, we provide illustrations that show more details on how the state and action are built.
Figure A.1a shows how to build the state representation and Figure A.1b how to compute the action representation of a particular region.
[Figure]
B E XTENDED EXPERIMENTAL SETUP
The segmentation network f is an adaptation of feature pyramid network (Lin et al., 2017) for semantic segmentation (similar to the segmentation branch of (Kirillov et al., 2019)), with a ResNet50 backbone (He et al., 2016), pretrained on ImageNet (Deng et al., 2009).
The network is pretrained on the full train set of a large-scale synthetic dataset, GTAV (Richter et al., 2016), therefore not requiring much human labelling effort.
Moreover, this dataset has the advantage of possessing the same categories as real datasets we experiment with.
The query network π , depicted in Figure B.1, is composed of two paths, one to compute state features and another to compute action features, fusing them at the end.
Each of the layers are composed of
Batch Normalization, ReLU activation and a fully-connected layer.
The state path and action path are composed of 4 and 3 layers, respectively, with a ﬁnal layer that fuses them together to get the global features; these are gated with a sigmoid, controlled by the KL distance distributions in the action representation.
The weights are updated at each step of the active learning loop, by sampling batches of 16 experience tuples from an experience replay buffer, sized 600 and 3200 for Camvid and Cityscapes, respectively.
We train both networks with stochastic gradient descent (SGD) with momentum.
We use the same learning rate for both the segmentation and query networks; 10 − 4 and 10 − 3 for Cityscapes and
Camvid respectively.
Weight decay is set to 10 − 4 for the segmentation network and 10 − 3 for the query network.
We used a training batch size of 32 for Camvid and 16 for Cityscapes.
[Figure]
Flatten
[Figure]
C C LASS FREQUENCIES AND PERFORMANCE PER CLASS
We show in Figure C.1 a more detailed plot for the class frequencies of regions that each of the methods chooses for labeling.
As the entropy of the class distributions in Figure 3 show, our method picks more regions containing under-represented classes.
Specially, it asks labels for more Person,
Rider, Train, Motorcycle and Bicycle pixels.
We observe that our B baseline picks more than 50% of pixels for only 3 classes that are overrepresented or have a medium representation: Building, Vegetation and Sky.
This could explain why the performance is worse than the H baseline.
Moreover, in table C.1, we extend Table 1 by adding the standard deviation for each result.
D Q UALITATIVE RESULTS
We compare our method qualitatively with the baselines in Figure D.1.
Baseline U asks for random patches.
Our method tends to pick more regions with the under-represented classes and small objects.
For instance, in the ﬁrst image in the left, our method asks for several regions of a Train, that almost has no samples in the training data.
In the second image, it focuses on Person, Bicycle and Poles.
In the third image, it asks for labels of the trafﬁc lights and a pedestrian on a bicycle.
Baselines B and H select some of those relevant regions, but miss a lot of them.
E A BLATION S TUDIES
In this section, we provide an ablation study on the state and action representation, the effect of labeling small regions versus full images, and the comparison of taking different regions per step.
[Table]
E. S TATE AND A CTION R EPRESENTATION
Here, we analyze the incremental effect of our design choices for the state and action representation on Cityscapes.
We use 3 pooling operations – min, average, max – to compress the entropy map of the region and use it in the state and action representation.
Also, KL divergences are added to the latter.
As seen in Table E.1, using only the max-pooled entropy map ( Ours - 1H ), the performance is slightly worse than H .
When we combine the information of the 3 pooled entropy maps ( Ours3H ), we outperform H baseline.
Moreover, when adding the two distribution of KL distances to our action representation ( Ours - 3H + KL ): between possible regions to label and the labeled set and between the region and the unlabeled set, we further increase the performance, getting our best state and action representations.
U
B
H
Ours
[Table]
E. R EGION VS . FULL IMAGE ANNOTATION
In this subsection
[Table]
we analyze the effect of asking for labels in regions instead of full images and the effect of the number of regions per step.
We compare the validation IoU when asking for pixel-wise labels for entire images versus pixel-wise labels for small regions.
In the ﬁrst case, we ask for one image at each step and, for the latter, we ask for 24 regions per step (pixel-wise, equivalent to one image).
As it is shown in Table E.2, asking for entire image labels has similar performance for all methods, that resemble Uniform performance when asking for region labels.
This indicates that, in order to select more informative samples, it is useful to split the images into patches (crops) and be able
E. I NFLUENCE OF STEP REGIONS
Empirically, our selector network is quite robust to the number of regions per step, as seen in Table E.3.
Therefore, we select 24 regions for CamVid, the one that yielded best results.
This is more efﬁcient to train than taking one region per step.
[Table] [Table]