L EARNING TO A NNOTATE P ART S EGMENTATION WITH G RADIENT M ATCHING
A BSTRACT
The success of state -of the -art deep neural networks heavily relies on the presence of large-scale labelled datasets, which are extremely expensive and time-consuming to annotate.
This paper focuses on tackling semi-supervised part segmentation tasks by generating high-quality images with a pre-trained GAN and labelling the generated images with an automatic annotator.
In particular, we formulate the annotator learning as a learning-to-learn problem.
Given a pre-trained GAN, the annotator learns to label object parts in a set of randomly generated images such that a part segmentation model trained on these synthetic images with automatic labels obtains superior performance evaluated on a small set of manually labelled images.
We further reduce this nested-loop optimization problem to a simple gradient matching problem and efﬁciently solve it with an iterative algorithm.
We show that our method can learn annotators from a broad range of labelled images including real images, generated images and even rendered images. Our method is evaluated with semi-supervised part segmentation tasks and signiﬁcantly outperforms other semi-supervised competitors when the amount of labelled examples is extremely limited
I NTRODUCTION
In recent years, deep neural networks have shown a remarkable ability to learn complex visual concepts from large quantities of labeled data.
However, collecting manual annotations for a growing body of visual concepts remains an extremely expensive practice, especially for pixelwise prediction tasks such as semantic segmentation.
A common and effective way of reducing the dependence on manual labels is ﬁrst learning representations from unlabeled data (Pathak et al., 2016; Noroozi &
Favaro, 2016; Gidaris et al., 2018; He et al., 2020) and then deploying learned representations to a supervised learning task (Chen et al., 2020; Zhai et al., 2019).
The insight of this paradigm is that some of the transferred representations are informative, easily identiﬁed, and disentangled, such that small amount of labeled data sufﬁces to train the target task.
A promising direction for learning such representations is using generative models (Kingma &
Welling, 2014; Goodfellow et al., 2014), since they attain the capability of synthesizing photorealistic images and disentangling factors of variation in dataset.
Recent methods (Goetschalckx et al., 2019; Shen et al., 2020; Karras et al., 2019; Plumerault et al., 2020; Jahanian et al., 2020; Voynov & Babenko, 2020; Spingarn-Eliezer et al., 2021; H¨ark¨onen et al., 2020) show that certain factors such as object shape and position in the images synthesized by generative adversarial networks (GANs) (Goodfellow et al., 2014) can be individually controlled by manipulating latent features.
Building on the success of powerful GAN models, Zhang et al.
(2021) and Tritrong et al.
(2021) demonstrate that feature maps in StyleGAN (Karras et al., 2019; 2020b;a) can be mapped to semantic segmentation masks or keypoint heatmaps through a shallow decoder, denoted as annotator in this paper.
Remarkably, the annotator can be trained after few synthesized images are labeled, and the joint generator and trained annotator can be used as an inﬁnite labeled data generator.
Nonetheless, labeling synthesized images with “human in the loop” has two shortcomings: (i) it can be difﬁcult to label the generated images of low quality; (ii) it requires continual human effort whenever the generator is changed.
To mitigate these limitations, a natural idea is to learn annotator from a collection of image-label pairs.
This is non-trivial because annotator only receives as input generator features which are lacked in labeled image datasets.
One straightforward way to address this issue is to inverse the generation process to obtain generator features for given images.
However, inversion methods can not ensure exact recovery of the generator features, leading to degraded performance of annotator (see Section C.1).
Another way to address this issue is to align the joint distribution of generated images and labels with that of real ones using adversarial learning as in SemanticGAN (Li et al., 2021).
However, adversarial learning is notoriously unstable (Brock et al., 2018) and requires large amount of data to prevent discriminator from overﬁtting (Karras et al., 2020a).
In this paper, we formulate the annotator learning as a learning-to-learn problem – the annotator learns to label a set of randomly generated images such that a segmentation network trained on these automatically labeled images obtains superior performance evaluated on a small set of manually labeled images.
This problem involves solving a nested-loop optimization, where the inner loop optimizes the segmentation network and the outer loop optimizes the annotator based on the solution of inner-loop optimization.
Instead of directly using end-to-end gradient-based meta learning techniques (Li et al., 2019; Pham et al., 2021), we reduce this complex and expensive optimization problem into a simple gradient matching problem and efﬁciently solve it with an iterative algorithm.
We show that our method obtains performance comparable to the supervised learning counterparts (Zhang et al., 2021; Tritrong et al., 2021) yet overcomes their shortcomings such that labeled data within a broader range including real data, synthetic data, and even out-of-domain data can be utilized.
Our method requires a large quantity of unlabeled data for pre-training GANs and a relatively small amount of labeled data for training annotator, which suits a semi-supervised learning setting.
Furthermore, considering that state-of-the-art unconditional GANs only produce appealing results on single-class images but struggle to model complex scenes, we evaluate our method on several part segmentation datasets of particular classes such as human face, horse, aeroplane, car, and cat.
Our contribution can be summarized as follows.
(i)
We formulate the learning of annotations for GANgenerated images as a learning-to-learn problem and propose an algorithm based on gradient matching to solve it.
As a consequence, a broad range of labeled data including real data, synthetic data, and even out-of-domain data is applicable.
(ii) We empirically show that our method signiﬁcantly outperforms other semi-supervised segmentation methods under the circumstances of extremely limited labeled data.
R ELATED W ORK
Semi-supervised learning Semi-supervised learning (SSL) augments the training of neural networks from a small amount of labeled data with a large-scale unlabeled dataset.
A rich body of work regularize networks on unlabeled data with consistency regularization.
In particular, networks are required to make consistent prediction over previous training iterations (Laine & Aila, 2017;
Tarvainen & Valpola, 2017; Izmailov et al., 2018), noisy inputs (Miyato et al., 2018), or augmented inputs (Berthelot et al., 2019b;a; Sohn et al., 2020).
This paradigm is also interpreted as a teacherand-student framework where teacher produces pseudo labels for unlabelled data to supervise the training of students.
Pseudo labels can be enhanced with heuristics (Lee et al., 2013; Berthelot et al., 2019a) or optimized towards better generalization (Pham et al., 2021;
Li et al., 2019).
Apart from image classiﬁcation tasks, the above ideas are also successfully applied to semantic segmentation tasks (Hung et al., 2018; Ouali et al., 2020; Ke et al., 2020; French et al., 2020).
Our work can also be interpreted from teacher-and-student perspective.
In contrast to the above methods that adopt homogeneous network structures for teacher and student, our work employ generative models as teacher and discriminative models as student.
Other SSL approaches explore the utility of unlabelled data from representation learning perspective.
In particular, deep neural networks are ﬁrst optimized on the unlabeled data with self-supervised learning tasks (Pathak et al., 2016; Noroozi & Favaro, 2016;
Gidaris et al., 2018;
He et al., 2020) to gain powerful and versatile representations, and then ﬁnetuned on the labeled data to perform target tasks (Chen et al., 2020; Zhai et al., 2019).
In addition, efﬁcient and transferable representations can also be learned with generative models that are typically trained without labels.
To this end, one promising direction is to decode generative features into dense annotations such as segmentation masks and keypoint heatmap to facilitate SSL, as done in SemanticGAN (Li et al., 2021), DatasetGAN (Zhang et al., 2021), RepurposeGAN (Tritrong et al., 2021) and our work.
However, our work differs from these works in the way that we formulate the annotator learning as a learning-to-learn problem which is solved with gradient matching.
Particualrly, unlike SemanticGAN (Li et al., 2021) that learns annotator and data generator jointly, our method presumes a ﬁxed pre-trained GAN and learns annotator only, which circumvents the complication of joint learning.
Semantic Part Segmentation Semantic part segmentation provides dense form of representations for both rigid and non-rigid objects, which is of great signiﬁcance in tremendous computer-vision tasks such as human parsing (Dong et al., 2014; Nie et al., 2018; Fang et al., 2018;
Gong et al., 2017), pose estimation (Yang & Ramanan, 2011; Shotton et al., 2011; Xia et al., 2017) and 3D object understanding (Yi et al., 2016; Song et al., 2017).
The challenge of part segmentation arises from viewpoint changes, occlusion, and lack of clear boundaries for certain parts.
Existing datasets manage to alleviate the problem by enrich the data and annotations.
CelebAMask-HQ (Lee et al., 2020a) is a large-scale face image dataset annotated with masks of face components.
Chen et al.
(2014) provides additional part segmentation upon PASCAL VOC dataset (Everingham et al., 2010).
Those datasets meet the demands of latest data-driven and deep neural network based method (Lee et al., 2020b; Wang et al., 2015; Wang & Yuille, 2015).
However, a well-curated part segmentation dataset containing ﬁne-grained pixel-wise annotation usually requires time-consuming manual labeling work, which holds back the development of this ﬁeld.
This problem also motivates the latest work (Zhang et al., 2021;
Li et al., 2021; Tritrong et al., 2021) to leverage a generative model as an inﬁnite data generator in a semi-supervised learning setting, which gain promising results.
Gradient matching Gradient matching is an important technique in meta learning (Li et al., 2018;
Sariyildiz & Cinbis, 2019), which is shown to be effective in domain generalization (Li et al., 2018), zero-shot classiﬁcation (Sariyildiz & Cinbis, 2019), and dataset condensation (Zhao et al., 2021; Zhao & Bilen, 2021) etc .
Although the general principle and gradient matching loss function in our work resemble those in the above works, we are motivated to solve a totally different problem.
In particular,
Gradient Matching Network (GMN) (Sariyildiz & Cinbis, 2019) employs gradient matching to train a conditional generative model towards application on zero-shot classiﬁcation.
Dataset Condensation (DC)
(Zhao et al., 2021; Zhao & Bilen, 2021) condenses a large-scale dataset into a small synthetic one with gradient matching.
Our work shows that gradient matching is also an effective way to learn annotations of GAN-generated images from limited labeled examples.
Moreover, in contrast to DC where the synthesized images are typically not realistic at all, our work are concerned about photo-realistic synthetic images.
M ETHOD 3. P RELIMINARY
Let G denote a generative model trained from a dataset D tot without labels.
G can be trained via adversarial learning (Goodfellow et al., 2014) such that it can generate photo-realistic images.
The generation process is formulated as
[Equation]
where z denotes a random variable drawn from a prior distribution P z which is typically a normal distribution, x denotes the generated image, and h denotes the hidden features in G .
An annotator
A ω is constructed to decode h into labels
[Equation]
where ω parameterizes the annotator.
Let D l denote a manually labeled dataset.
Given a ﬁxed G , our problem is to learn annotator A ω from D l such that A ω can annotate images generated from G .
[Figure]
𝑮 𝑨 𝝎 𝜵 𝜽 𝓛 𝒈 (𝜽,𝝎) 𝜵 𝜽 𝓛 𝒍 (𝜽) 𝓛 𝒈 (𝜽,𝝎)
𝓛 𝒍 (𝜽)
L a [Figure]
d
D a t a 𝑺 𝜽 𝑺
[Figure]
Revisiting supervised annotator learning DatasetGAN (Zhang et al., 2021) and RepurposeGAN (Tritrong et al., 2021) train annotators in a supervised learning manner.
First, a set of synthetic images { x i } Ni =1 generated from G are selected and annotated by humans and their generator features { h i } Ni =1 are reserved.
These features and manual labels { y i } Ni =1 constitute the labeled dataset D l = { ( x i , h i , y i ) } Ni =1 .
Second, an annotator is trained by minimizing the loss computed between annotator prediction and ground truth labels,
[Equation]
[Equation] where f denotes per-pixel cross-entropy function for segmentation tasks.
This method, however, requires D l to be a synthetic and generator-speciﬁc set, which has the following two defects.
(i) Details might be illegible in synthetic images, which increases the difﬁculty of annotation.
(ii) Data needs to be re-labeled whenever the generator is changed.
Although projecting labeled images into GAN latent space can alleviate the above drawbacks, it leads to degenerate performance due to inexact inversion (see Section C.1 in the appendix).
3. L EARNING TO ANNOTATE WITH G RADIENT M ATCHING
Learning to learn We formulate the problem of learning A ω as a learning-to-learn problem like Li et al.
Pham et al.
In particular, we introduce a segmentation network S θ which is a deep neural network parameterized with θ .
It takes as input an image and outputs a segmentation mask.
This segmentation network learns the target task only from the automatically labeled generated images, which is formulated as
[Equation]
where f denotes the loss function for the target task.
The optimal solution of Equation 4 depends on the ω , which is therefore denoted as θ ∗ ( ω ) .
The performance of learned segmentation network can be evaluated with loss on the labeled dataset D l as
[Equation]
The aim of learning the annotator is to produce high-quality labels such that the learned segmentation network achieve minimal loss on D l , which is formulated as
[Equation]
This is a nested-loop optimization problem, where the inner loop optimizes segmentation network (Equation 4) and the outer loop optimizes annotator based on the solution of the inner loop.
This optimization problem does not require manual labeling of generated images.
To this end, any labeled images, including real labeled images, can serve as D l .
While this problem can be solved with MAML-like (Finn et al., 2017) approaches (Li et al., 2019; Pham et al., 2021), its solution involves expensive unrolling of the computational graph over multiple θ states, which is memory and computation intensive.
Hence, we reduce it into a simple gradient matching problem as follows.
See
Section C.1 in the appendix for further comparison and discussion.
Gradient matching Inspired by Li et al.
(2018), we reduce the nested-loop optimization problem (Equation 6) into a gradient matching problem.
First, as commonly done in gradient-based meta learning approaches, the optimization of student (Equation 4) is approximated with one-step gradient descent,
[Equation]
where η denotes the learning rate and θ 0 denotes the initialized parameter.
By plugging this equation into L l ( θ ∗ ( ω )) , unfolding it as Taylor series and omitting the higher order term, we have
[Equation]
The last approximation is conditioned on (cid:107) η ∇ θ L g ( θ 0 , ω ) (cid:107) ≤ (cid:15) , where (cid:15) is a very small amount.
By further omitting the constant term L l ( θ 0 ) in Equation 8, we obtain the approximation of Equationas
[Equation]
We notice a connection between this constrained optimization problem with the following unconstrained optimization problem
[Equation]
where to minimize this objective, one needs to simultaneously maximize the dot product between the gradients and minimize the norm of ∇ θ L g ( θ 0 , ω ) .
Note that − ∇
[Table]
[Equation]
[Table]
[Equation]
[Table]
, and the constant term and constant coefﬁcient can be omitted without changing the optimal solution.
We re-write the learning objective of Equation 10 as
[Equation]
which computes the cosine distance between the gradients of segmentation network on automatically labeled generated data and manually labeled data (see Section B in appendix for dealing with gradient matching in multi-layer neural networks)
To this end, the nested-loop optimization (Equation 6) is reduced into a simple one that minimizes the gradient matching loss (Equation 11).
Illustration of this procedure is presented in Figure 1.
Algorithm Note that the gradient matching problem learns an annotator for a particular network parameterized with θ 0 .
However, a desirable annotator should produce universal automatic labels for training any segmentation network with previously unseen parameters.
Hence, we present an alternate training algorithm of annotator and segmentation network in Algorithm 1, where the segmentation network is iteratively updated using the automatically labeled synthetic data and the annotator is iteratively updated with gradient matching.
This algorithm allows the annotator to be optimized over various segmentation network states, and therefore leads to a segmentation-network-agnostic annotator.
The gradient matching problem for each segmentation network parameter is solved with K -step gradient descent and typically K = 1 works well in practice.
Implementation details are available in the appendix (see Section B).
E XPERIMENTS 4. S ETUP
Datasets We evaluate our method on six part segmentation datasets: CelebA , Pascal-Horse , PascalAeroplane , Car-20 , Cat-16 , and Face-34 .
CelebA
(Liu et al., 2015) is a large-scale human face dataset where 30,000 images are annotated with up to 19 part classes, also known as CelebAMaskHQ (Lee et al., 2020a).
We consider a subset of 8 face classes and use the ﬁrst 28,000 images as unlabeled set, the last 500 images as test set, and the rest 1500 images as labeled set, as in Li et al.
Pascal-Horse and Pascal-Aeroplane are constructed by taking images of horse and aeroplane from PascalPart (Chen et al., 2014) which provides detailed part segmentation annotations for images
[Table]
sample a batch of labeled examples
[Table] [Table]
[Table] [Table]
// compute gradients
[Table] [Table]
[Table]
[Table]
// sample a batch of synthetic data
[Table] [Table]
[Table] [Table] [Table]
[Table]
// compute gradients
[Table]
[Table]
// compute gradient matching loss
[Table]
[Table] [Table] [Table]
// update annotator parameters
[Table]
[Table] [Table] [Table]
[Table]
// sample a batch of synthetic data
[Table]
[Table] [Table]
synthetic
[Table] [Table]
[Table] [Table]
from Pascal VOC 2010 (Everingham et al.) The selected images are cropped according to bounding box annotations.
Full process details are available in the appendix.
We ﬁnally obtain 180 training images, 34 validation images, and 225 test images to constitute Pascal-Horse, and 180 training images, 78 validation images, and 269 test images to constitute Pascal-Aeroplane.
In Pascal-Horse and Pascal-Aeroplane, additional images from LSUN (Yu et al., 2015) are utilized as unlabeled set.
Car-20, Cat-16, and Face-34, released by Zhang et al. (2021), are three part-segmentation datasets annotated with 20, 16, and 34 part classes for car, cat, and human face, respectively. We use the same datasets to make a fair comparison to DatasetGAN (Zhang et al., 2021).
In these datasets, all the training images are images generated from pre-trained StyleGAN (Karras et al., 2019) while all the test images are real images.
Car-20 contains 16 training images and 10 test images; Cat-16 contains 30 training images and 20 test images; Face-34 contains 16 training images and 20 test images.
Pre-trained generators We use models of StyleGAN family (Karras et al., 2019; 2020b;a) that are either trained by ourselves or publicly available.
Details are provided in the appendix (see Section B).
Evaluation We use mean intersection over union (mIoU) to evaluate the performance of segmentation networks.
On CelebA, Pascal-Horse, and Pascal-Aeroplane, we use validation set to select checkpoints and report mIoU across all foreground classes, denoted as “FG-mIoU”, on test set.
On
Car-20, Cat-16 and Face-34, we follow the setting of DatasetGAN (Zhang et al., 2021) and report the cross validation mIoU across all classes including background on test set.
4. S EMI -S UPERVISED P ART S EGMENTATION
Since our method has mild requirement for the labeled data, in this section, we show the effectiveness of our method under three different circumstances: real images as labeled data, synthetic images as labeled data, and out-of-domain images as labeled data.
Real images as labeled data Following SemanticGAN (Li et al., 2021), our methods are compared against SSL methods that have codes available 2 , including Mean Teacher (MT) (Tarvainen & Valpola,
[Figure] [Figure]
G
S
[Figure]
Face
[Figure]
# annotations: 16 # classes:
[Figure]
Car-
# annotations: 16 # classes:
TL † -
DeepLabv 21.58 ± 0.
45.77 ± 1.
33.91 ± 0.
SSL (Mittal et al., 2019) † -
DeepLabv 24.85 ± 0.
48.17 ± 0.
44.51 ± 0.
DatasetGAN (Zhang et al., 2021)
StyleGAN
DeepLabv3 ‡ 32.63 ± 0.
54.55 ± 0.
67.53 ± 2.
StyleGAN
U-Net (cid:93) 31.36 ± 0.
53.84 ± 0.
[Table]
Ours
StyleGAN
DeepLabv 33.89 ± 0.
52.58 ± 0.
63.55 ± 2.
StyleGAN
U-Net 32.64 ± 0.
[Table]
69 ±
[Table]
[Table]
StyleGAN
DeepLabv 33.56 ± 0.
55.10 ±
[Table]
61.21 ±
[Table]
StyleGAN
U-
[Table]
2017; Li et al., 2020), Guided Collaborative Training (GCT) (Ke et al., 2020), an adversariallearning-based semi-supervised segmentation method (AdvSSL) (Hung et al., 2018), as well as
SemanticGAN (Li et al., 2021).
We benchmark performances on CelebA, Pascal-Horse
[Table]
and PascalAeroplane with respect to different amount of labeled data and present the results in Figure 2.
Qualitative results are available in the appendix (see Figure D.7, D.8, D.9).
Our method signiﬁcantly outperforms other SSL methods when labeled data is extremely limited.
When the number of labeled images decreases, the performance of our method drops mildly while performances of other methods decreases drastically.
We attribute it to our utilization of the highly interpretable generator features.
Notably, SemanticGAN (Li et al., 2021) also exploits the generative features to produce segmentation masks but its performance degrades faster than our methods as the number of labeled images decreases.
We conjecture it is due to that the adversarial learning employed by SemanticGAN typically requires sufﬁcient data to prevent discriminators from overﬁtting.
The edges of our method become marginal and even unclear under a case of relatively large amount of labeled data.
It can be foreseen that our method is less competitive in the large-scale setting. One possible reason is that the quality of pre-trained GAN needs to be further improved, which requires more research in the future. Synthetic images as labeled data We comprehensively compare our method to a supervised learning counterpart, DatasetGAN (Zhang et al., 2021) on Car-20, Cat-16, and Face-34.
We follow
[Figure]
the same setting as in DatasetGAN, using exactly the same synthetic images as training data.
Quantitative results are presented in Table 1 and qualitative comparison is available in the appendix (see Figure D.12, D.11, D.10).
First, our method achieves the same performance level with DatasetGAN on Face-34 and Cat-16.
Second, as a merit of our method, we can upgrade the generator to StyleGAN2 (Karras et al., 2020b;a) without requiring extra human effort.
Upgrading generator, however, brings modest improvement on Face-34 and Cat-16 and even a bit decrease on Car-20.
This suggests that synthesis quality measured by metrics like FID can not precisely depict the disentanglement of generator features.
Finally, our method could adapt to different segmentation architectures such as DeepLabv3 (Chen et al., 2017) and U-Net (Ronneberger et al., 2015).
We further use our trained annotator to generate synthetic segmentation datasets for training downstream segmentation networks.
Table 1 shows the results when different source networks, which are used for training annotator with gradient matching, and different downstream networks are used.
First, our learned synthetic datasets can be used to train downstream segmentation networks of different architecture that achieve good performances.
This indicates that our learned synthetic dataset is not architecture-speciﬁc.
Second, DeepLabv3 as source network generally leads to higher performances than U-Net does, suggesting that in our method the segmentation networks used for gradient matching has an impact on annotator learning.
Out-of-domain images as labeled data We further consider a challenge where out-of-domain images constitute labeled data.
In particular, we render a set of car images from annotated 3D CAD models provided by CGPart 3 (Liu et al., 2021) using graphic tools ( e.g . blender).
Since 3D models are annotated with parts, the ground truth part segmentation of these rendered images are available via 3D projection.
These rendered images vary in viewpoints and textures and have signiﬁcant domain gap compared to realistic car images due to inexact 3D models, unrealistic textures, unnatural lighting etc ., raising challenges for learning annotators.
Despite so, as presented in Figure 3, our method still learns reasonable annotators that produce quality segmentation labels for GAN generated images.
It is noteworthy that it is nearly impossible for DatasetGAN (Zhang et al., 2021) and RepurposeGAN (Tritrong et al., 2021) to utilize such out-of-domain labeled data.
4. A BLATION S TUDY
The frequency of updating student The results of an ablation study w.r.t.
the interval of updating student network, K , is shown in Figure 4.
Experiments are run on Face-34 and Cat-16 with
DeepLabv3 as segmentation network and StyleGAN as generator such that pairs of generator features and ground truth segmentation are available to quantitatively measure the performance of annotator.
It shows that the annotator learns signiﬁcantly faster when the segmentation network is updated more frequently.
From the visualization, it can be seen that very detailed parts, e.g . pupils in Faceand eyes in Cat-16, emerge sooner if segmentation network is updated more frequently.
It can be concluded that it is more effective to match gradients for various student states than optimize gradient matching for single student state.
Gradient matching on partial student parameters As the implementation of gradient matching requires maintenance of computation graph of backpropagation through student networks, one may [-Generated Image # Step ( 𝟏𝟎 𝟑 )    
-]concern the training efﬁciency for very deep neural networks.
Here we present a simple strategy to
[Figure]
Step ( ×𝟏𝟎 𝟑 )   20  𝟒 , 𝑲 = 𝟏𝟔 , 𝟒 , 𝑲 = 𝟏𝟔 ,
[Figure]
[Figure] [Table] [Figure]
[Table]
[Figure]
[Table]
trade off the training efﬁciency and overall performance: gradient matching can be done only on part of student parameters.
An ablation study on Cat-16 using StyleGAN2 as generator and DeepLabv as student network is done and the results are shown in Table 2.
We consider network blocks in
DeepLabv3 with ResNet-101 as backbone from bottom to top: Res1 , Res2 , Res3 , Res4 and
Res5 that are residual blocks in the backbone, and Head that is the segmentation head.
Results show that by skipping the bottom blocks for gradient matching, the training speed is improved while the segmentation performance is affected little.
We conjecture the reasons can be two folds: (i) the backbone parameters are pre-trained such that little needs to be changed during training; (ii) the gradients accumulate more randomness as they backpropagate to the bottom blocks of deep neural networks.
Therefore, in practice, we only match the gradients of Head in DeepLabv3 by default.
Additional results are available in the appendix (see Section C.3).
C ONCLUSION
We propose a gradient-matching-based method to learn annotator which is able to label generated images with part segmentation by decoding the generator features into segmentation masks.
Unlike existing methods that require labeling generated images for training an annotator, our method allows usage of a broader range of labeled data including realistic images, synthetic images, and even rendered images from 3D CAD models.
On the benchmark of semi-supervised part segmentation, our method signiﬁcantly outperform other semi-supervised segmentation methods under the circumstances of extremely limited labeled data but becomes less competitive under large-scale setting which requires feature research.
The effectiveness of our method is validated on a variety of single-class datasets including well-aligned images, e.g . CelebA and Face-34, as well as images in the wild part of which contain cluttered backgrounds, e.g . Pascal-Horse, Pascal-Aeroplane, Car-20 and Cat-16.
In terms of more complex scenes, discussion and investigation are further needed.
With the rapid progress of generative modeling, it is promising to use more powerful generative models and explore more computer vision problems under different challenging settings in the future. A CKNOWLEDGEMENTS
This work was supported by the National Key R&D Program of China under Grant 2018AAA0102801,
National Natural Science Foundation of China under Grant 61620106005, and Beijing Municipal
Science and Technology Commission grant Z201100005820005. Hakan Bilen is supported by the EPSRC programme grant Visual AI EP/T028572/1. R EPRODUCIBILITY S TATEMENT
We provide all the details about the datasets, network structure and training curriculum in the appendix.
Code to reproduce our main results is publicly available at https://github.com/ yangyu12/lagm .
A D ATASET D ETAILS
CelebA Following Li et al.
(2021), we only consider 8 part classes: background, ear, eye, eyebrow, skin, hair, mouth, and nose.
The protocol for merging the 19 part classes into 8 classes is presented in
Table A.1(Left).
CelebAMask-HQ contains 30,000 annotated images in total.
We split this dataset into unlabeled set, training set, validation set and test set as in Table A.1(Right).
[Table]
Pascal-Horse & Pascal-Aeroplane Pascal Part (Chen et al., 2014) provides part segmentation annotations of 20 object classes for images in Pascal VOC 2010.
Following other work on part segmentation (Tsogkas et al., 2015; Zhao et al., 2019; Tritrong et al., 2021), we merge the ﬁne-grained part classes into 6 classes for horse and aeroplane.
The merging protocol is presented in Table A.2.
We crop the image patches that contain the object of interest according to bounding box annotations.
Concretely, bounding boxes that have IoU greater than 0.05 with other boxes or empty annotations are discarded to ensure single object appears in single image patch.
The patches that have any side less a certain number of pixels (32 for horse; 50 for aeroplane) are also abandoned.
The crop regions are extended from bounding boxes to square boxes and the regions outside images are padded with zero.
These processed patches are ﬁnally resized to 256 × 256 resolution to serve the training and evaluation process.
For horse and aeroplane, We further split the patches in ofﬁcial VOCtrain split in training and validation set, and take the patches in ofﬁcial VOC 2010 val split as test set.
This procedure ﬁnally provides 180 images as training set ( i.e . labeled set), 33 images as validation set, and 223 images as test set in Pascal-Horse, and 180 images as labeled set, 78 images as validation set and 266 images as test set in Pascal-Aeroplane.
As unlabeled images, we center-crop the 200,000 images in the large-scale LSUN (Yu et al., 2015) archive and resize them to 256 ×resolution for both horse and aeroplane.
[Table]
CGPart As a part segmentation dataset composed of 3D CAD models from 5 vehicle categories with 3D part manual annotations, CGPart (Liu et al., 2021) manage to render a large amount of image dataset with part mask labels by adopting blender, a well-known graphical software.
In this paper, we select category car as training and evaluation dataset and make moderate modiﬁcations upon the rendering pipeline.
Parts are merged into 15 classes, where excessively ﬁne-grained parts that seldom emerge are combined together to avoid sample imbalance.
Part merging protocol is presented in Table
A.3.
Viewpoints are sampled around the object with azimuth ranged in full 360 ◦ .
Elevation angle is varied from 0 ◦ to 80 ◦ , yet higher sampling probability is assigned to interval [0 ◦ , 30 ◦ ] .
The distance from object to viewpoint is carefully set so that the object is located at center and full in frame.
Random color is painted upon the object and random texture sampled from COCO dataset (Lin et al., 2014) and Pixar-One-Twenty-Eight (Sisson, 2018) is rendered on the ground and the surroundings.
4 , 000 images of 512 × 384 are rendered as full train dataset.
When training, only 1 , 000 of them in total are randomly sampled as ofﬁcial trainset.
The original val split of CGPart is adopted as test dataset, consisting of 40 images.
Test masks are merged under the same merging protocol in inference and evaluation phases.
[Table]
B I MPLEMENTATION D ETAILS
Gradient matching loss Following DC
(Zhao et al., 2021), the distance between gradients is measured with cosine similarity.
The computation of gradient matching is quite similar as in DC.
Here, we re-state this procedure for clarity.
In particular, considering a multi-layer neural network S θ that is parameterized with θ , the gradient matching loss is computed as an average over layerwise losses as D ( ∇ θ L l , ∇ θ L g )
=
1 L (cid:80) Li =
1 d ( ∇ θ ( i ) L l , ∇ θ
( i ) L g ) , where i denotes the layer index, L denotes the number of layers for gradient matching and
[Equation]
12) where A j and B j are ﬂattened gradient vectors for each neural node j .
The minor difference of our practice compared DC is that we average the layerwise gradient distances whereas DC sum them.
It is easier to handle numerical issues by using our practice in our problem.
In terms of gradient matching for normalization layers such as Batch Normalization (BN), we follow the practice of DC to ignore the gradient matching of learnable parameters in BN.
The BN layer is set as train mode during gradient matching.
We ﬁnd it works well.
Pre-trained GAN We use StyleGAN family (Karras et al., 2019; 2020b;a) as our pre-trained
GANs.
For each dataset, We either train GANs by ourselves or use publicly available models.
The training conﬁgurations or links of the pre-trained GANs are listed in Table B.4.
For GANs trained by ourselves, we use the checkpoints with lowest FID in the historical ones.
Annotator structure Annotator is a neural network that takes as input the generator features and output segmentation masks.
DatasetGAN (Zhang et al., 2021) and RepurposeGAN (Tritrong et al., 2021)
ﬁrst upsample feature maps to full resolution (resolution of output image) and then concatenate
[Table]
these feature maps along channel.
This procedure results in extremely high-dimensional feature vectors for every pixel.
These features are fed into either a pixelwise multiple layer perceptrons (MLPs) as in DatasetGAN (Zhang et al., 2021) or a convolutional neural networks (CNNs) as in (Tritrong et al., 2021).
Nonetheless, this practice consumes a lot of GPU memory and run painfully slow even on high-end modern GPUs.
That is also why DatasetGAN cannot consume all image feature vectors in a batch during training.
Therefore [Table] to make training annotator more efﬁcient, as commonly done in segmentation or detection network architectures, we fuse the multi-scale feature maps with a feature pyramid network (FPN) (Lin et al., 2017) structure and decode the fused features into segmentation masks with consecutive convolutional layers.
This FPN structure saves a lot of GPU memories and computation, which allows us to consume multiple full images in one mini-batch during training.
An illustration of our annotator structure is presented in Figure B.1.
For generator features that are fed into annotator, we use outputs of all convolutional layers in StyleGAN as in DatasetGAN (Zhang et al., 2021) and use outputs of all synthesis blocks, each of which typically contains two consecutive convolutional layers, in StyleGAN2.
[Figure]
Segmentation network structures In our experiments, we employ two prevalent segmentation network structures: U-Net (Ronneberger et al., 2015) and DeepLabv3 (Chen et al., 2017) with ResNet-101 pretrained on ImageNet as backbone.
For U-Net, we reference the implementation from https://github.com/milesial/Pytorch-UNet and train it from scratch on target datasets.
For DeepLabv3, we use the the built-in implementation in PyTorch library 4 .
For all the experiments on CelebA, Pascal-Horse, and Pascal-Aeroplane, we report the results (Figure 2) using DeepLabv3 as segmentation networks.
Training details of our method We use pixelwise cross entropy loss as f ( · , · ) for training student network and computing gradients.
Both the annotator and the student are optimized with SGD optimizer with learning rate 0 .
001 and momentum 0 .
By default, we jointly train an annotator and a student with K = 1 and batch size 2 for 150 , 000 steps.
Training details of competing semi-supervised methods For MT, GCT and AdvSSL, we train the networks for 10 , 000 iterations with batch size 16 for labeled data and 8 for unlabeled data.
Learning rate is decayed with power 0 .
9 every iteration.
During training process, the input images are randomly resized to a scale in [160 , 640] and cropped to 256 × 256 .
For the rest of the hyper-parameters, we use the default settings as in PixelSSL.
C F URTHER A NALYSIS C. C OMPARISON TO O THER B ASELINES
Inversion baseline One way to make the supervised learning of annotator be able to utilize existing labeled images is to (i) project images into GAN latent space to acquire generator features, (ii) use the joint generator features and ground truth labels to train an annotator.
We call this method “inversion baseline”.
We evaluate the performance of this method on Cat-16, Face-34, and Carusing StyleGAN2 as generator and make a comparison to our method in Table C.5.
Concretely, we ﬁrst use the projection method provided by Karras et al.
(2020b) to compute the latent style codes that result in a reconstruction of labeled images.
Second, the latent style codes are forwarded to the generator to acquire generator features.
We present examples of such reconstruction in Figure C. which shows that the reconstruction is quite accurate.
However, the performance of inversion baseline fails to combat our method and the automatic labels for a randomly generated image are of low quality (see Figure C.3).
We argue that it is non-trivial to inverse the GAN generation process to acquire features for speciﬁc images.
Even though the reconstruction of images is satisfying, one is still unable to exactly recover the generator features, which leads to an degraded learning of annotator.
[Table]
Pseudo-labeling baseline Another way to employ existing labeled images to train an annotator in supervised learning manner is to use a trained segmentation model to predict pseudo labels for the generated images.
We refer to this method as “pseudo-labeling baseline”.
In particular, this baseline method involves three stages: (i) a segmentation model is trained on the labeled data; (ii) the trained segmentation model is used to produce pseudo labels for every random image generated from generator, which further supervise the learning of an annotator; (iii) the joint generator and trained annotator are used for generating automatically labeled data which is further used for training a downstream segmentation network.
As shown in Table C.6, pseudo-labeling baseline performs badly, especially on Car-20, which can be accounted for the following reasons.
Our method can be understood as distilling knowledge from a [Figure] generative model to a discriminative model, whereas this pseudo-labeling baseline is more like the other way around – distilling knowledge from a discriminative model to a generative model.
The highly interpretable generator features make our method possible to learn annotations from only a few labeled data.
In contrast, the discriminative model can hardly learn to acquire valuable knowledge from extremely limited labeled data. [Figure] [Table] [Figure]
[Table]
[Figure]
[Table] [Figure] [Table] Our method with MAML One direct solution to Equation 6 is to use MAML-like algorithm, which unrolls the inner loop with multi-step stochastic gradient descent, reserves the computation graph, computes the meta loss, and ﬁnally backpropagates the gradients for updating the annotator.
We refer to this baseline as “Ours- K -step-MAML” and summarize this algorithm as Algorithm 2.
Table C.7 shows the experimental comparison of gradient matching versus MAML, where we use η i = 0 . 1 in MAML for all experiments.
The following conclusions can be drawn from these results.
First, gradient matching generally leads to higher performance and obtains more consistent results across different datasets and different segmentation network architectures than MAML.
Second,
MAML probably requires multiple inner-loop steps (more than two) to match the performance of gradient matching, given that it fails to do so with two inner-loop steps.
This is very inefﬁcient and requires huge computation resources.
Furthermore, in practice, we observe that MAML appears unstable and sensitive to hyperparameters ( e.g . inner-loop learning rate) across different datasets and different network architectures, whereas gradient matching nearly requires neither dataset-speciﬁc nor network-architecture-speciﬁc hyperparameters. [Table]
[Table]
parameters for [Table] [Table] sample
[Table]
batch of
[Table]
// compute gradients on
[Table]
// update segmentor parameters
[Table]
// sample a batch of labeled examples
[Table]
// compute meta loss
[Table]
// update annotator parameters
[Table]
// sample a batch of synthetic data
[Table]
[Table]
loss on synthetic data
[Table]
// update segmentor parameters
[Table]
C. T HE I MPACT OF S TUDENT S TATES
As we present an alternate learning algorithm between annotator and student
[Table]
we study the impact of updating student states.
In particular, we investigate the training convergence and overall performance of annotator when different series of student states are employed for gradient matching.
We present in Figure C.4 & Figure C.5 the evolution of annotator performance evaluated on training set during training process, where “baseline” denotes the default setting and the following settings are further compared.
(i)
The student parameters are randomly re-initialized rather than learned from generated images, denoted as “random re-init”.
(ii)
The student is pretrained to acquire a fair performance ( ∼ 40 .
7% mIoU evaluated on test set), denoted as “pre-trained”.
(iii)
The student is a DeepLabv network that has backbone randomly initialized rather than pretrained on ImageNet, denoted as “scratch”.
(iv) The student parameters, either randomly initialized or pre-trained, are ﬁxed, denoted as “random init & ﬁxed” and “pre-trained & ﬁxed”, respectively.
First, “Fixed” students result in a gradient matching problem on single student state, which learns annotator with very poor performance compared to the default setting.
This result suggests that matching gradients for a series of student states is vital to learning annotator.
Second, even though it is crucial to require annotator to produce student-agnostic labels, a randomly re-initialization strategy does not work well (see results of “random re-init” v.s. “baseline”).
Moreover, the results of “pre-trained & ﬁxed” is better than “random init & ﬁxed”, which suggest that a well-performed student provides more informative gradients such that by gradient matching the annotator can be more effectively learned.
Third, the conclusion that a well-performed student provides more informative gradients can be further veriﬁed by the fact that matching gradients of pre-trained students accelerates the learning of annotator (see results of “baseline” v.s. “pre-trained” and “baseline” v.s. “scratch”).
[Figure] [Figure]
C. T RADE OFF T RAINING E FFICIENCY WITH P ERFORMANCE
Gradient matching on partial network parameters We present additional results on in Table C.8.
U-Net takes as input an image and ﬁrst produce feature map at multiple downsampled scales with consecutive convolutional and pooling layers.
These feature maps are then processed by consecutive convolutional and upsampling layers to be mapped back to original input scale.
Notably, skip connection is employed to connect the multi-scale feature maps along the downsampling path and upsampling path.
We consider the following particular group of layers in U-Net: (i) Up v.s.
Down , (ii) Input v.s. Output , and (iii) U1 v.s. U2 .
Up and Down denote the convolutional layers along the upsampling and downsampling path, respectively.
Input and Output denote the input convolutional layer and the output convolutional layer, respectively.
U1 and U2 denote the convolutional layers that process feature maps at 1 / 1 and 1 / 2 scales, respectively.
From the results we can see that matching gradients of layers near the output end is more effective than doing so at layers near the input end.
[Table]
Resolution of input to segmentation network Another way to improve the training efﬁciency is to reduce the resolution of input to student network.
Note that the resize operation can be implemented with interpolation which is differentiable.
We show in Figure C.6 and Table C.9 that by reducing the resolution, the training speed improves signiﬁcantly, but the annotator performance and downstream segmentation networks compromises.
Therefore, in practice, the student resolution can be tuned by trading off the training efﬁciency and effectiveness.
[Figure] [Figure]
[Table]
[Figure]
[Table]
C. T HE I MPACT OF G ENERATOR
As our segmentation network is trained only on the generated images (Algorithm 1 Line 9 ∼ 11) and a very shallow network is employed to convert generator features into annotations the segmentation performance is inevitably affected by the quality of pre-trained GANs.
which suggests that the quality of generative features signiﬁcantly affect the quality of annotations.
Table C.10 shows that premature
GANs – ones with high FID score – generally leads to lower segmentation performance.
[Table]
C. C OMPARISON TO D ATASET GAN IN O THER A SPECTS
Real images v.s. synthetic images as labeled data Since our method removes the necessity of labeling synthetic images, we study if replacing the labeled synthetic images with real ones could improve the segmentation performance.
Table C.11 shows an evaluation on CelebA-test at 512 ×resolution across 8 face classes.
DatasetGAN, which uses 16 annotated synthetic images for training, achieves 70.01 mIoU performance.
Our method achieves a bit higher performance when using the exactly same synthetic images as DatasetGAN.
Moreover, the performance of our method can be further improved with around 6.7% mIoU by replacing the labeled synthetic images with real ones in
CelebA-train.
This suggests that the capability of using real labeled data exhibits the superiority of our method over DatasetGAN.
[Table]
Our method with more labeled data on Car-20 Considering our method does not match the performance of DatasetGAN on Car-20, we further investigate if this gap can be narrowed with the help of more labeled data.
Table C.12 presents the performance of our method when more labeled data is used on Car-20.
It shows that with increased number of labeled data, the performance of our method gradually approaches that of DatasetGAN.
[Table]
D Q UALITATIVE R ESULTS
I npu t
G T
N u m b e r o f l a b e l e d exa m p l es
I npu t
G T
N u m b e r o f l a b e l e d exa m p l es
I npu t
G T
N u m b e r o f l a b e l e d exa m p l es
Synthetic Data Segmentation Results
I npu t
G T
D a t ase t
G AN
O u r s
D a t ase t G
AN
O u r s
[Figure]
I npu t
G T
D a t ase t
G AN
O u r s
D a t ase t G
AN
O u r s
[Figure]
I npu t
G T
D a t ase t
G AN
O u r s
D a t ase t G
AN
O u r s