Weakly-Supervised Semantic Segmentation via Sub-category Exploration

In plain words

Teaching a network to find objects using only image-level tags ("this photo contains a cat") is cheap, but the networks tend to light up only the most telling part of each object. The classifier already gets the right answer from that part, so nothing pushes it to look at the rest.

The paper adds a second, harder training task that does not replace the first. For each tag, the method groups the network's own image features into a few clusters with a standard clustering algorithm. Each cluster becomes a made-up sub-category label, and the network is trained to predict both the original tag and the sub-category. Because the features change during training, the clusters are recomputed and the process repeats for a few rounds.

The resulting highlight maps cover more of each object, which gives the rest of a segmentation pipeline a better starting point. The method changes only how the maps are learned, not how they are refined or how the final segmentation network is trained.

TL;DR

A method paper with code (MIT, first author's repository) for the initial localisation step of weakly-supervised segmentation on PASCAL VOC 2012. Start from a ResNet-38 classifier trained on the 20 parent tags, cluster each class's features into K=10 pseudo sub-categories, and train the parent and sub-category heads jointly with the sub-category loss weighted by 5 §3.2 §3.3. Repeat the clustering and training for 3 rounds §3.3. Use the parent-head CAM as the seed for the usual refinement and segmentation steps §3.1 §3.3. The paper reports better CAMs than AffinityNet's baseline and final segmentation results of 66.1 mIoU (val) and 65.9 (test) §4.4 tab. 4. Several training details are missing from the paper and differ from the released code (optimizer, batch size, epochs, starting point of each round); see Paper vs code before reimplementing.

Problem

With image-level labels, the usual approach is to train a classifier and read its class activation map (CAM) as a rough object location §1 §3.1. CAMs from a classifier highlight discriminative parts, because the classification loss is already satisfied by those parts §1 §3.2. Existing WSSS pipelines improve the refinement step and the final training, but they inherit the incomplete seeds from the first step §1 §2. The paper targets that first step and asks whether better features alone give more complete CAMs §1 §2.

Method

Overview §3.1. The parent network is a classifier trained with image-level labels Y_p and a multi-label loss \mathcal{L}_p §3.1. The CAM for class c is the parent classifier weight applied to the feature map, and it is normalised by its maximum §3.1 eq. 1. The method adds a sub-category head that shares the feature extractor E, and alternates two steps: cluster the features of each parent class, then train on the parent and sub-category labels §3.2 alg. 1. The final CAM still comes from the parent head §3.1 eq. 1.

Sub-category objective §3.2. For each parent class p_c there are K sub-categories. An image has a parent label Y_p^c and a sub-category label per parent class; if the parent label is 0, all its sub-category labels are 0 §3.2. A second head H_s with parameters \theta_s uses the same E as the parent head H_p §3.2. Both heads use the multi-label classification loss §3.2.

Sub-category discovery §3.2 eq. 2. For each parent class, the features f = E(I) of the images that contain that class are clustered into K groups. The cluster assignment of an image becomes its pseudo sub-category label §3.2.

\min_{D \in \mathbb{R}^{d \times k}} \frac{1}{N^c} \sum_{i=1}^{N^c} \min_{Y_s^c} \lVert f - T Y_s^c \rVert_2^2, \quad \text{s.t.}\; Y_s^{c\top} 1_k = 1 eq. 2

Joint training §3.2 eq. 3. With the pseudo labels Y_s in hand, the parent and sub-category losses are summed with weight \lambda on the sub-category term §3.2. The paper sets \lambda = 5 §3.3.

\min_{\theta_p, \theta_s} \frac{1}{N} \sum_{i=1}^{N} \mathcal{L}_p(H_p(f_i), Y_p) + \lambda\, \mathcal{L}_s(H_s(f_i), Y_s) eq. 3

Iterative procedure §3.2 alg. 1. The algorithm first optimises \{E, H_p\} on Y_p alg. 1. It then repeats: extract features, cluster each parent class with eq. 2, and optimise \{E, H_p, H_s\} with eq. 3 alg. 1. After the loop, the CAM is computed with eq. 1 alg. 1. The paper says the iteration gives the clustering a feature space that was itself shaped by the sub-category objective §3.2.

# Iterative sub-category exploration for the CAM  [alg. 1] [§3.2]
# Input:  images I, parent labels Y_p (N, C), K, lambda, number of rounds
# Output: trained extractor E and parent head H_p; CAM via [eq. 1]

def sub_category_cam(images, parent_labels, K, lam, rounds):
    # Parent-only training of E and H_p on the image tags  [alg. 1]
    extractor = train_parent_only(images, parent_labels)
    for round_nb in range(rounds):
        # Features of every training image from the current extractor  [§3.2]
        features = extract_features(extractor, images)      # (N, D)
        # Per parent class, cluster the features of the images that have
        # that tag; the cluster index is the pseudo sub-label  [eq. 2]
        sub_labels = cluster_per_class(features, parent_labels, K)
        # Train E, H_p and H_s jointly on the tags and pseudo labels  [eq. 3]
        extractor = train_joint(images, parent_labels, sub_labels, lam)
    # Parent-head CAM is the output used by the rest of the pipeline  [eq. 1]
    return extractor

Segmentation pipeline §3.3. The CAM is turned into pixel pseudo labels with the random-walk affinity refinement of AffinityNet, and a dense CRF is applied as common practice §3.3. The pseudo labels train a DeepLab-v2 model with a ResNet-101 backbone §3.3. The paper states that the method is not tied to a particular refinement or segmentation training method §3.1.

Network and training as described §3.3. A wide ResNet-38 backbone with 38 convolution layers is followed by a 3×3 convolution with 512 channels, global average pooling, and two fully connected heads, one for the parent classes and one for the sub-categories §3.3 fig. 2. The backbone is pretrained on ImageNet and fine-tuned on PASCAL VOC 2012 §3.3. Augmentation uses horizontal flip, random crop, colour jitter and random scaling §3.3. The model is implemented in PyTorch on a single Titan X with 12 GB §3.3.

Cost §3.3 §5. The paper says the method adds no extra complexity or inference time to the model §5. The extra cost is in training: the classifier is trained once per round, with feature extraction and clustering before each round §3.2 alg. 1.

Hyperparameters

Stated values are from the paper; "code" values are from the released repository at code_commit.

Results

Author-reported. Setup: PASCAL VOC 2012, 21 classes including background, metric mIoU §4.1. Table 1 scores the CAM and CAM after random-walk refinement (CAM+RW); the setup of the AffinityNet row is not stated §4.2.

Iterative improvement on the validation set (CAM, ResNet-38). Round #0 is the original CAM tab. 2 §4.3.

Sensitivity to K: the text says the method is consistently better than the CAM (K = 1) for K = 5, 10, 20 and 50, but the values are only in the plot §4.3 fig. 4. K = 10 was chosen §4.3.

Segmentation on PASCAL VOC 2012 (ResNet-101 DeepLab-v2, pseudo labels from the method) §4.4. Per-class results on val tab. 3 give these mIoU values; the caption splits them into with and without CRF, but the source table does not split the baseline rows cleanly.

Comparison with other methods on val and test (backbone and whether they improve the initial response, marked "Init. Res.") tab. 4:

Figures carry the qualitative claims: CAMs cover more of the object than the baseline CAM fig. 3; clusters on the last round group objects by size, context, type, pose and interaction with other classes fig. 5 §4.3; t-SNE of the classifier weights places a sub-category near the parent class it co-occurs with fig. 6; and the final segmentation results are shown on val fig. 7.

Paper vs code

Official code: the first author's repository Juliachang/SC-CAM is linked in the paper §3.3 and its LICENSE is MIT (copyright line names the author). The repository contains the CAM network, feature extraction, pseudo-label generation, training and inference. It does not contain the AffinityNet refinement or the DeepLab-v2 segmentation network; the README points to them README.md.

1. Optimiser and learning rate. Paper: Adam, initial learning rate 1e-3, weight decay 5e-4 §3.3. Code: SGD with momentum 0.9 and polynomial decay with power 0.9 torchutils.py; the base learning rate passed by every script call is 0.01 script.sh script.sh script.sh; the layers get multipliers 1, 2, 10 and 20 train_cls.py.

2. Starting point of each round. Paper: the iteration updates the feature extractor and both heads §3.2 alg. 1. Code: every round starts from the ImageNet weights, because the script passes the same pretrained_model to all three training calls script.sh script.sh script.sh script.sh. The round-1 features come from AffinityNet's res38_cls.pth, not from a parent-only model trained in this repository script.sh README.md. The parent-only branch exists in the code resnet38_cls.py but no script calls it.

3. Checkpoints. Paper: "converges after 3 rounds" §3.3. Code: the rounds use the epoch-30, epoch-40 and epoch-50 checkpoints out of a 61-epoch run that saves every 10 epochs script.sh train_cls.py train_cls.py. The paper does not name the epoch.

4. Feature used for clustering. Paper: f = E(I) in \mathbb{R}^D eq. 2. Code: the 4096-dimensional pooled feature, without normalisation, clustered by scikit-learn KMeans with random_state=0 and max_iter=300 extract_feature.py create_pseudo_label.py. The default number of K-means initialisations is set by the scikit-learn version, which the repository does not pin.

5. Pseudo labels for images with several objects. Paper: each image has a sub-category label per parent class, with Y_s^{c,k} \in \{0,1\}^k §3.2. Code: merges the per-object cluster labels into one 200-dimensional multi-hot vector per image create_pseudo_label.py, and discards an image's merged label when two of its objects fall into the same cluster create_pseudo_label.py. The paper does not mention this filter. The filtered list is shorter than the unfiltered one, yet later steps index the filtered labels with positions from the unfiltered list create_pseudo_label.py create_pseudo_label.py create_pseudo_label.py. This appears to misalign labels and images; verify before reuse.

6. Head width. Paper: 3×3 convolution with 512 channels, then pooling and two fully connected heads §3.3. Code: the last block outputs 4096 channels through a dilated bottleneck with dropout 0.5 resnet38d.py resnet38_cls.py. The two heads are 1×1 convolutions on the 4096-channel pooled feature resnet38_cls.py.

7. Frozen layers. Paper: not stated. Code: conv1a and the b2 blocks are frozen, and every BatchNorm layer is in eval mode with its parameters frozen resnet38_cls.py resnet38_cls.py resnet38d.py.

8. Loss. Paper: "standard multi-label classification loss" §3.2, not written out. Code: multilabel_soft_margin_loss on the pre-sigmoid outputs, and the total loss is the parent loss plus 5 times the sub-category loss train_cls.py.

9. CAM at inference. Paper: M^c = \theta_p^{c\top} f, normalised by its maximum eq. 1 §3.1. Code: each class map is multiplied by the image's ground-truth image-level tags, passed through ReLU, summed over 4 scales and their flips, and divided by the maximum plus 1e-5 infer_cls.py infer_cls.py resnet38_cls.py infer_cls.py. The tag masking and multi-scale averaging are not stated in the paper.

10. Sub-category CAM. Code: the 200-channel CAM is normalised per parent class infer_utils.py and merged to parent classes by taking a max over sub-categories infer_utils.py. Only the 20-class CAM is saved by the released script infer_cls.py. The paper uses the parent-head CAM §3.1.

11. Background threshold. The released evaluation turns CAMs into label maps with a constant background score of 0.15 infer_cls.py infer_utils.py. The paper gives no threshold for pseudo labels.

12. Training-set size. Paper: 10,528 augmented images §4.1. Repository: voc12/train_aug.txt has 10,582 lines train_aug.txt, and the code comments also use 10,582 create_pseudo_label.py.

No discrepancy was found for K = 10, λ = 5 and 3 rounds §3.3 train_cls.py script.sh.

Pitfalls

• Eq. 3 omits the backbone. The equation minimises over \theta_p, \theta_s only, but the text and alg. 1 update E as well §3.2 eq. 3 alg. 1. The code updates all parameters, apart from the frozen layers in item 7 of Paper vs code train_cls.py.
• Symbol clash in eq. 2. D is used for both the feature dimension and the matrix variable, and T is called the centroid matrix. Read it as K-means over the N^c features of class c, with K centroids §3.2 eq. 2.
• Parent-only start is not specified. alg. 1 begins with a parent-only training step, but the paper does not say how long it runs or what it starts from. The code takes it from AffinityNet's weights script.sh.
• Each round is a full retrain. The code creates new heads and reloads the ImageNet backbone each round script.sh resnet38_cls.py. Cluster identities do not carry over between rounds, and training cost is roughly three full trainings. If you warm-start instead, you change the method; measure it.
• Feature scale matters for K-means. The clustered features are raw pooled activations with no normalisation extract_feature.py. The paper does not state a normalisation, so changing one changes the clusters.
• Pseudo labels depend on the ground-truth tags. At CAM generation the released code masks each class with the image's ground-truth tags infer_cls.py infer_cls.py, as is usual for WSSS. Do not reuse these CAMs for images without tags.
• Table 4's "Ours" row. The 66.1 val figure matches the with-CRF row of tab. 3, but the paper does not say outright that tab. 4 includes CRF §4.4.
• Test numbers are not reproducible here. Test results come from the official evaluation website §4.1, and the refinement and segmentation steps are external to this repository README.md.
• Figure-only sensitivity. The K sweep is only in fig. 4. Read the values from the figure yourself before you cite them.
• Augmentation values are missing from the paper. The random long side of 256 to 512 and the 224 crop train_cls.py train_cls.py are not stated in the paper, and they change the features that K-means sees.

Lineage

• Builds on: the CAM classifier of Zhou et al. §2 §3.1; the training recipe and the random-walk refinement of AffinityNet (Ahn and Kwak) §3.1 §3.3 README.md; the ResNet-38 backbone of Wu et al. §3.3; DeepLab-v2 with a ResNet-101 backbone for segmentation §3.3; and the alternation between clustering and training from DeepCluster (Caron et al.) §2. The paper gives no arXiv identifiers for these works, so none are linked here.
• Compared against: SEC, MCOF, DSRG, DCSP, SeeNet, FickleNet, CIAN, BDSSW, OAA and Zeng et al. §2 tab. 4. Among them, SeeNet, OAA and FickleNet also target the initial response tab. 4.
• Superseded by: not known from content in this wiki.