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 their predicted labels obtains low segmentation error on a small validation 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 analytically 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 labelled data.
However, collecting manual annotations for a growing body of visual concepts remains a costly practice, especially for pixel-wise 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 transferring 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 a small amount of labelled 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 the 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 a few synthesized images are manually labelled, and the joint generator and trained annotator can be used as an inﬁnite labelled data generator.
Nonetheless, labelling 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.
One straightforward way to mitigate these limitations is to inverse the generation process to obtain generator features for given images.
In this way, annotator, which only receives generator features as input, can be trained from existing labelled data, not necessarily generated data.
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 a large amount of data to prevent discriminators 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 labelled images obtains low prediction error on a small validation set of manually labelled 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 training annotators can utilize labelled data within a broader range , including real data, synthetic data, and even out-of-domain data
Our method requires a large quantity of unlabeled data for pre-training GANs and a relatively small amount of labelled data for training annotators, which drops into a semi-supervised learning setting.
Our method relies on unconditional GANs, of which the performance limit the application of our approach.
As the state-of-the-art GANs only produce appealing results on single-class images but struggle to model complex scenes, we focus on part segmentation of particular classes and avoid multi-object scenes.
Our contribution can be summarized as follows.
(i)
We formulate the learning of annotations for
GAN-generated images as a learning-to-learn problem and propose an algorithm based on gradient matching to solve it.
Consequently, a broad range of labelled 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 in the few-label regime.
R ELATED W ORK
Semi-supervised learning Semi-supervised learning (SSL) (Zhu, 2005) augments the training of neural networks from a small amount of labelled data with a large-scale unlabeled dataset.
A rich body of work regularizes networks on unlabeled data with consistency regularization.
In particular, networks are required to make consistent predictions 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 teacher-and-student framework where the 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 a teacher-and-student perspective.
In contrast to the above methods that adopt homogeneous network structures for teacher and student, our work employs generative models as a teacher and discriminative models as a student.
Other SSL approaches explore the utility of unlabelled data from a 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 labelled data to perform target tasks (Chen et al., 2020; Zhai et al., 2019).
In addition, generative models can also learn efﬁcient and transferable representations 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 that we formulate the annotator learning as a learning-to-learn problem that is solved with gradient matching.
Particularly, 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, which decomposes rigid or non-rigid objects into several parts, 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 a lack of clear boundaries for certain parts.
Existing datasets manage to alleviate the problem by enriching 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 labelling 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 gains 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 our work’s general principle and gradient matching loss function resemble those in the above works, we are motivated to solve a 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 labelled examples.
Moreover, in contrast to DC, where the synthesized images are typically not realistic, our work is 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 ( h
[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 represents 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]
𝑮 𝑨
[Figure]
𝜵 𝜽 𝓛 𝒈 (𝜽,𝝎) 𝜵 𝜽 𝓛 𝒍 (𝜽) 𝓛 𝒈 (𝜽,𝝎)
𝓛 𝒍 (𝜽)
[Figure]
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-labelled whenever the generator is changed.
Although projecting labelled images into
GAN latent space can alleviate the above drawbacks, it leads to degraded 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 θ , 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 problem is a nested-loop optimization problem, where the inner loop optimizes the segmentation network (Equation 4) and the outer loop optimizes the annotator based on the solution of the inner loop.
This optimization problem does not require manual labelling of generated images.
To this end, any labelled images, including real labelled 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 segmentation network (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).
An 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 labelled synthetic data, and the annotator is iteratively updated with gradient matching.
This algorithm allows the annotator to be optimized over various segmentation network states, leading 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 an unlabeled set, the last 500 images as a test set, and the rest 1500 images as a labelled 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
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
sample a batch of labeled examples
[Table]
[Table]
// compute gradients
[Table]
[Table]
[Table]
// sample a batch of synthetic data
[Table]
[Table]
// compute gradients
[Table]
[Table]
[Table]
[Table]
// compute gradient matching loss
[Table]
// update annotator parameters
[Table]
[Table] [Table]
[Table]
// sample a batch of synthetic data
[Table]
[Table]
[Table] [Table]
synthetic
[Table]
[Table] [Table]
for images 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 obtaintraining images, 34 validation images, and 225 test images to constitute Pascal-Horse, 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 an 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 the validation set to select checkpoints and report mIoU across all foreground classes, denoted as “FG-mIoU”, on the 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 the test set.
4. S EMI -S UPERVISED P ART S EGMENTATION
Since our method has a mild requirement for the labelled data, in this section, we show the effectiveness of our method under three different circumstances: real images as labelled data, synthetic images as labelled data, and out-of-domain images as labelled 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 a different amount of labelled 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 labelled data is extremely limited.
When the number of labelled images decreases, the performance of our method drops mildly while the performances of other methods decrease 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 labelled 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 under the case of a large number of labels.
One possible reason is that the quality of pre-trained GAN needs to be further improved, which requires more research in the future.
(see
Section C.5 for more discussion)
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 precisely 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.
However, upgrading the generator brings modest improvement on Face-34 and Cat-16 and even a bit decrease on Car-20.
This result 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 we use different source networks for training annotators with gradient matching and different downstream networks.
First, our learned synthetic datasets can be used to train downstream segmentation networks of different architecture that achieve good performances.
This result indicates that our learned synthetic dataset is not architecture-speciﬁc.
Second, DeepLabv3 as a source network generally leads to higher performances than U-Net, suggesting that the segmentation networks used for gradient matching affect annotator learning in our method.
Out-of-domain images as labeled data We further consider a challenge where out-of-domain images constitute labelled 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 a signiﬁcant domain gap compared to realistic car images due to inexact 3D models, unrealistic textures, artiﬁcial 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 labelled data.
4. A BLATION S TUDY
The frequency of updating segmentation network The results of an ablation study w.r.t.
the interval of updating segmentation network, K , is shown in Figure 4.
Experiments are run on Faceand Cat-16 with DeepLabv3 as segmentation network and StyleGAN as generator such that pairs of generator features and ground truth segmentation are available to measure the performance of the annotator quantitatively.
It shows that the annotator learns signiﬁcantly faster when the segmentation network is updated more frequently.
The visualization shows that very detailed parts, e.g . pupils in Face-34 and eyes in Cat-16, emerge sooner if the segmentation network is updated more frequently.
It is more effective to match gradients for various segmentation network parameters than optimize gradient matching for single segmentation network parameters.
Gradient matching on partial segmentation network parameters As the implementation of gradient matching requires maintenance of computation graph of backpropagation through segmentation networks, one may be concerned about the training efﬁciency for very deep neural networks.
Here we
[Figure]
×𝟏𝟎 𝟑 )   20  𝟒 , 𝑲 = 𝟏𝟔 , 𝟒 , 𝑲 = 𝟏𝟔 ,
[Figure]
[Figure] [Table] [Figure]
[Table]
[Figure]
[Table] [Figure] [Figure] [Table]
[Figure]
[Table]
present a simple strategy to trade off the training efﬁciency and overall performance: gradient matching can be done only on the part of segmentation network parameters.
Table 2 presents an ablation study about matching blocks on Cat-16, where StyleGAN2 is the generator and DeepLabv3-R101 is the segmentation network.
We consider network blocks in DeepLabv3-R101 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 skipping the bottom blocks for gradient matching improves the training speed 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 can label generated images with part segmentation by decoding the generator features into segmentation masks.
Unlike existing methods that require labelling generated images for training an annotator, our method allows a broader range of labelled 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 outperforms other semi-supervised segmentation methods under the circumstances of extremely limited labelled data but becomes less competitive under large-scale settings, which requires future 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 modelling, 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
The appendix provides all the details about the datasets, network structure, and training curriculum.
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 segmentation network and computing gradients.
Both the annotator and the segmentation network are optimized with SGD optimizer with learning rate 0 .
001 and momentum 0 .
By default, we jointly train an annotator and a segmentation network 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]
// copy segmentor parameters for simulated SGD
[Table]
// sample a batch of synthetic data
[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]
// compute loss on synthetic data
[Table]
// update segmentor parameters
[Table]
T HE I MPACT OF S EGMENTATION N ETWORK S TATES
As we present an alternate learning algorithm between annotator and segmentation network
[Table]
we study the impact of updating segmentation network states, i.e . segmentation parameters.
In particular, we investigate the training convergence and overall performance of annotator when different series of segmentation network 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 segmentation network parameters are randomly re-initialized rather than learned from generated images, denoted as “random re-init”.
(ii)
The segmentation network is pretrained to acquire a fair performance ( ∼ 40 .
7% mIoU evaluated on test set), denoted as “pre-trained”.
(iii)
The segmentation network is a DeepLabv3 network that has backbone randomly initialized rather than pretrained on ImageNet, denoted as “scratch”.
(iv) The segmentation network 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 segmentation network state, which learns annotator with very poor performance compared to the default setting.
This result suggests that matching gradients for a series of segmentation network states is vital to learning annotator.
Second, even though it is crucial to require annotator to produce segmentation-networkagnostic 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 segmentation network provides more informative gradients such that by gradient matching the annotator can be more effectively learned.
Third, the conclusion that a well-performed segmentation network 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]
(%)) is evaluated on training set of Car-20.
Segmentation network is U-Net.
[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 segmentation 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 input resolution to segmentation network can be tuned by trading off the training efﬁciency and effectiveness.
[Figure]
)) evaluated on training and validataion set during training process on Face-34.
[Table]
[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. D ISCUSSION ON L ARGE -S CALE S ETTING
As shown in Fig.
[Table]
the performance of our method begins to saturate and even beaten by other semi-supervised segmentation method when the number of labels grow large.
To investigate the performance of our method under large-scale setting
[Table]
we further run experiments with more unlabeled and labeled data and present the results in Table C.11.
It shows that the performance of our method even slightly drops when the number of unlabeled data or the number of labels increases and is signiﬁcantly beaten by supervised learning.
We hypothesize the reason to be that our method is seriously limited by the performance of GANs which still struggle to ﬁt large-scale datasets.
In contrast, under large-scale setting, supervised learning becomes stronger since more labels are available.
How to address this issue requires further research in the future
[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.12 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.13 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
[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
[Figure]