ImageNet-21K Pretraining for the Masses

In plain words

Most vision models are first trained on ImageNet-1K, a 1,000-class subset of the much larger ImageNet-21K. The bigger set has about 21,000 categories, but it is awkward to use. It has no official validation split, it is over a terabyte of images, and its labels are messy. Each picture gets one label, but that label can sit at any level of the WordNet word hierarchy. The same cow can be tagged "cow" in one photo and "animal" in another, which confuses an ordinary classifier.

The authors first clean the dataset: they drop very rare categories, hold out a fixed validation split and shrink the images. Then they use the hierarchy as extra training signal. Instead of one classifier over all categories, they train one classifier per hierarchy level, and each image trains every level for which it has a label. The levels are weighted so that the crowded lower levels do not dominate. The authors report that this pretraining beats ImageNet-1K pretraining for every model they tried, including small mobile networks, and beats the official ImageNet-21K weights for ViT and Mixer.

TL;DR

Preprocess ImageNet-21K into a fixed, cleaned set of 11,221 classes at 224 px §2.1. Pretrain with semantic softmax: one softmax per WordNet depth, with each image labelled at every depth its ancestors reach, and depth losses reweighted §3.3 eq. 2 eq. 3. Author-reported gains over ImageNet-1K pretraining hold on all 30 model-dataset cells tested tab. 4. The released code omits the semantic KD training and its depth weights differ from the paper, so reproduce the paper's numbers with care [Paper vs code].

Problem

ImageNet-21K is large and public, but it is rarely used for pretraining. The paper names four obstacles: each image carries a single label that may sit at any WordNet level, so the same object can be labelled inconsistently; there is no official train/validation split; training takes longer than ImageNet-1K; and the raw data is about 1.3 TB §1 §2.1 fig. 2. Earlier large-scale ImageNet-21K pretraining work targeted large models, or used private data such as JFT-300M §1 §5.1.

Method

Preprocessing §2.1

1. Remove classes with fewer than 500 images. The paper reports 21,841 classes reduced to 11,221, and 14,197,122 images reduced to 12,358,688 (about 13% of images removed) §2.1.
2. Hold out 50 images per remaining class as a standard validation split §2.1.
3. Resize every image to 224 x 224 pixels, ignoring aspect ratio ("squish" resizing). This reduces the data from 1.3 TB to 250 GB §2.1 app. B.

The processed set is called ImageNet-21K-P §2.1.

Semantic labels §2.2

Each class is a WordNet synset. Following its hypernym links gives the class a depth (its number of ancestors). An image labelled "swan" also carries the labels for its ancestors at each depth, here "animal, vertebrate, bird, aquatic bird, swan" for depths 0 to 4 §2.2. The processed set has 11 depths §2.2. The paper says most classes sit at the lower depths §2.2 app. A.

Baselines §3.1 §3.2

• Single-label: one softmax over all classes with cross-entropy and label smoothing §3.1 app. B. The paper initialises from ImageNet-1K weights and trains 80 epochs. Random initialisation needs 140 epochs §3.1.
• Multi-label: one sigmoid per class, with the binary losses summed over all N = 11,221 classes eq. 1 §3.2. The paper uses asymmetric loss (ASL) §3.2 app. B and reports it beats cross-entropy and focal loss app. E.

Semantic softmax §3.3

• Use one softmax per depth k (0 to 10). Each softmax only sees the logits of classes at its own depth §3.3.
• An image labelled at depth d has a valid target at depths 0 to d. Only those softmax layers receive gradient for that image §3.3 fig. 3.
• The total loss is the sum of the per-depth losses, weighted by W_k eq. 3. The paper defines W_k as 1/O_k, with O_k given by eq. 2. Eq. 2 and the code disagree; see Paper vs code.

L_{tot} = \sum_{n=1}^{N} L(\sigma(z_n), y_n) eq. 1

O_k = \sum_{j=0}^{k-1} N_j, \quad N_j = \text{number of classes at depth } j eq. 2

L_{tot} = \sum_{k=0}^{K-1} W_k L_k, \quad W_k = \frac{1}{O_k} eq. 3

# Semantic softmax loss for one batch  [§3.3] [eq. 3]
# Input:  logits (B, C) for all C classes; labels (B,) class index
# Output: scalar loss
def semantic_softmax_loss(logits, labels):
    # Split the class logits into one group per depth  [§3.3]
    depth_logits = split_by_depth(logits)  # K groups, (B, C_k)
    total = 0
    for depth in range(K):
        # Label of each image at this depth, or -1 if it is shallower  [§2.2]
        depth_labels = ancestor_at_depth(labels, depth)  # (B,)
        active = depth_labels >= 0
        # Cross-entropy, label smoothing 0.2; inactive rows add 0  [app. B]
        loss = smoothed_cross_entropy(
            depth_logits[depth], depth_labels, active)
        # Weight the depth with W_k; see Paper vs code  [eq. 3]
        total = total + weights[depth] * loss
    return total

Semantic knowledge distillation §3.4

The paper adds a distillation term to semantic softmax. For each depth, a student's and a teacher's per-depth probability distributions are compared with a distance; the paper uses mean squared error, reporting that it converged faster than KL divergence §3.4 eq. 4. Each depth's term is weighted by a teacher confidence P_i, which is the total probability of the teacher's top 5% classes at that depth §3.4 app. F eq. 5. The paper says P_i is set to 1 when the image's ground-truth depth is above i app. F. The teacher model is not named in the paper, and the weight between this term and the classification loss is not stated.

L_{KD_i} = \mathrm{KDLoss}(T_i, S_i) eq. 4

L_{KD} = \sum_{i=0}^{K-1} P_i \, L_{KD_i} eq. 5

# Semantic KD loss for one batch  [§3.4] [eq. 4] [eq. 5]
# Input:  student_logits, teacher_logits (B, C) for all C classes
# Output: scalar loss
def semantic_kd_loss(student_logits, teacher_logits):
    # Split both logit sets into one group per depth  [§3.4]
    student_groups = split_by_depth(student_logits)
    teacher_groups = split_by_depth(teacher_logits)
    total = 0
    for depth in range(K):
        # Softmax over the classes of this depth only  [§3.4]
        student_probs = softmax(student_groups[depth])
        teacher_probs = softmax(teacher_groups[depth])
        # Teacher confidence: mass of its top 5% classes  [app. F]
        confidence = top_five_percent_mass(teacher_probs)  # (B,)
        # Squared error per class, summed, weighted per image  [app. K]
        distance = squared_error(student_probs, teacher_probs)
        total = total + sum(confidence[:, None] * distance)
    return total

Hyperparameters

Results

Author-reported. Setup: ImageNet-21K-P built from the Fall11 release, pretrained with the schemes above, then transferred to downstream datasets: ImageNet-1K, iNaturalist 2019, CIFAR-100, Food-251, MS-COCO, Pascal-VOC and Kinetics-200 §4.2 app. D. Baselines are ImageNet-1K pretraining (weights from timm) and the official ImageNet-21K weights for ViT-B-16 and Mixer-B-16 §5.1 §5.2. Table 2's architecture is not stated in its caption; the appendix says TResNet-M unless stated otherwise app. B.

Pretraining schemes, single-label vs multi-label vs semantic softmax tab. 2:

Semantic softmax beats both baselines on six of seven datasets. Multi-label is better on CIFAR-100 §4.2. The KD column improves on plain semantic softmax on every dataset §4.2 tab. 2.

ImageNet-21K-P vs official ImageNet-21K weights tab. 3:

All 14 cells favour the paper's pretraining §5.1 tab. 3.

ImageNet-1K pretraining vs ImageNet-21K-P (top-1 % or mAP %, ImageNet-1K value first, then ours) tab. 4:

All 30 cells favour ImageNet-21K-P §5.2 tab. 4.

ImageNet-1K transfer with ImageNet-21K-P pretraining plus KD tab. 9: MobileNetV3 78.0, OFA-595 81.0, ResNet50 82.0, Mixer-B-16 82.2, TResNet-M 83.1, TResNet-L 83.9, ViT-B-16 84.4 (top-1 %, input 224). The authors state these are new state-of-the-art results at that resolution §5.3 tab. 9.

Supporting results (appendix)

• Single-label upstream accuracy on ImageNet-21K-P, top-1 / top-5 %: MobileNetV3 37.8 / 66.3; ResNet50 42.2 / 72.0; TResNet-M 45.3 / 75.2; TResNet-L 45.5 / 75.6 tab. 5.
• Multi-label loss comparison, micro / macro mAP %: cross-entropy 47.3 / 73.9; focal 47.4 / 74.1; ASL 48.5 / 74.7 tab. 6.
• Multi-label pretraining, cross-entropy vs ASL (ImageNet-1K / iNaturalist / Food-251 / CIFAR-100 / MS-COCO / Pascal-VOC / Kinetics-200): 79.6 vs 81.0; 69.4 vs 71.0; 74.3 vs 75.2; 89.9 vs 90.6; 79.1 vs 80.6; 87.6 vs 87.9; 81.1 vs 81.9 tab. 7.
• Semantic KD vs single-label KD (same column order as tab. 2): 81.5 vs 82.2; 72.4 vs 72.7; 76.0 vs 76.1; 91.0 vs 91.7; 81.6 vs 82.2; 89.0 vs 89.8; 83.6 vs 84.4 tab. 8.
• Open Images pretraining (multi-label) vs ImageNet-21K-P (semantic softmax), on the same seven datasets, Open Images first: 81.0 vs 81.4; 70.7 vs 72.0; 74.8 vs 75.8; 89.4 vs 90.4; 80.5 vs 81.3; 89.6 vs 89.7; 81.6 vs 83.0 tab. 12.
• MS-COCO detection, mAP %: 42.9 (ImageNet-1K) vs 44.3 (ImageNet-21K) tab. 13. INRIA Holidays retrieval, mAP %: 81.1 vs 82.1 tab. 14.
• Fall11 vs Winter21 release of ImageNet-21K-P, downstream top-1 / mAP %: ImageNet-1K 81.4 vs 81.2; iNaturalist 72.0 vs 71.8; Food-251 75.8 vs 75.5; CIFAR-100 90.4 vs 90.5; MS-COCO 81.3 vs 81.1; Pascal-VOC 89.7 vs 90.1; Kinetics-200 83.0 vs 82.8 tab. 11. The paper's main numbers use Fall11 app. I.

Paper vs code

The official repository is Alibaba-MIIL/ImageNet21K, linked from the abstract and whose README names the first author. Cloned at HEAD 6e4c6a4, licence MIT, copyright Alibaba-MIIL LICENSE. The repository contains the training scripts, the semantic-softmax and KD code, and the preprocessing scripts. The WordNet semantic tree file is not in the repository and must be downloaded README.md.

• Depth weights (eq. 2, eq. 3). Paper: O_k is the sum of N_j over j < k, and W_k = 1/O_k eq. 2 eq. 3. Code: W_k = S_0 / S_k, where S_k is the sum of class counts at depth k and deeper, and the value is capped at 20 semantics.py semantics.py. Applied as a weighted sum semantic_loss.py. The paper's own argument (a sample activates every depth at or below its own) implies a sum over j at or above k, which is what the code does. Eq. 2 as printed is not what the code implements.
• Semantic depths. Paper: 11 depths §2.2. Code: heads only for depths with more than one class semantics.py. The normalisation still counts all depths semantics.py.
• Loss averaging. Paper: loss is zero where a depth has no label, with no normalisation rule stated §3.3. Code: the masked cross-entropy is averaged over all B rows, so inactive rows dilute the mean semantic_loss.py.
• Label smoothing. Paper: 0.2 stated for single-label training only app. B. Code: 0.2 applied per depth, spread over that depth's classes semantic_loss.py.
• Semantic KD is not in the training scripts. The KD loss is implemented in a helper module, but no training script calls it semantics.py. The README lists "KD training code" under "To be added" README.md. The paper's KD results therefore cannot be reproduced from the release.
• Teacher confidence. Paper: P_i = 1 when the image's ground-truth depth is above i; otherwise the top-5% mass app. F. Code: only the top-5% mass, computed from the teacher's predictions alone, with no ground-truth check semantics.py.
• Teacher model. Not named in the paper §3.4. The code takes a teacher's logits as an input and does not say what the teacher is.
• KD loss scale. Paper pseudocode uses a sum over classes and batch app. K. Code also sums, not averages semantics.py. The paper gives no weight between KD and classification losses.
• Class threshold. Paper: remove classes with fewer than 500 labels §2.1. Script: removes classes with 500 or fewer processing_script.sh.
• Validation split. Paper: 50 per class, no selection rule §2.1. Script: first 50 files in directory order, deliberately not shuffled "for reproducibility" processing_script.sh.
• Resizing. Paper: all images resized to 224 §2.1. Script: images whose longer side is below 224 are copied unchanged resize.py. The data loader then resizes every image to 224 x 224 at load time data_loader.py.
• Cutout. Paper gives factor 0.5 only app. B. Code fills the box with a random colour and sets its side to half the image side augmentations.py.
• Optimiser. Paper: Adam, lr 3e-4, one-cycle, true weight decay app. B. Code matches, with weight decay 1e-4 by default train_semantic_softmax.py create_optimizer.py.
• Initialisation. Paper: from ImageNet-1K weights §3.1. Code: loads a weights file given by --model_path factory.py. The default path is a local file that the user must supply train_semantic_softmax.py.
• Downstream training. Paper: full settings for seven datasets app. D. Repository: the only transfer example is a timm command for CIFAR-100 with Mixer-B-16 Transfer_learning.md. Other downstream settings are not in the repository.
• Multi-label training. The paper's multi-label scripts are not in the repository. The single-label script uses the same optimiser and schedule with cross-entropy plus smoothing train_single_label.py.

Pitfalls

• Eq. 2 and the code disagree on the depth weights. Eq. 2 sums over j < k, and W_k = 1/O_k is undefined at k = 0 because O_0 is an empty sum. The code uses a suffix sum, normalises so W_0 = 1 and caps weights at 20. Implement what the code does, and note the choice.
• The KD recipe is not reproducible from the release. The teacher, the mix with classification loss, the KD learning schedule and the P_i ground-truth rule are not all available. Treat semantic KD results as author-reported without a runnable baseline.
• Teacher confidence rule is ambiguous. "Higher than i" could mean strictly above the ground-truth depth or at or above it. The paper's own example (an image at depth 5 is relevant to depths 0 to 5) suggests "at or above". The code ignores the ground truth altogether.
• Masked loss denominator. Inactive depth rows count as zero in the mean. This makes the effective step for sparse depths smaller than the per-depth loss suggests. The paper does not state the denominator.
• Label smoothing for semantic softmax is not in the paper. The code uses 0.2 per depth. If you change this, expect different numbers.
• Validation split is deterministic but arbitrary. The first 50 files in directory order are used. Papers that compare against these numbers must use the same selection or they are not comparable.
• Totals in §2.1 do not add up. The paper says 12,358,688 images in total, but its train count plus test count is 12,358,684 app. I. The four-image gap is small, but check your own counts before you claim a match.
• Depth tree is external. The WordNet depth tree is downloaded, not built by the repository, and the repository has no script that builds it. Use the Fall11 tree for the paper's numbers. The Winter21 release has a different class set and its own tree README.md processing_instructions.md app. I.
• Dataset terms. The repository states that use of ImageNet-21K-P is subject to image-net.org terms of access processing_instructions.md. The paper does not state the dataset licence separately. Check the terms before redistributing any derived data.
• Compute. The paper estimates 54 hours for 80 epochs on 8 V100 GPUs with ImageNet-1K initialisation §3.1. It says a single GPU would take two to three weeks app. L.
• Results are author-reported and use fixed hyperparameters. The README says the article's comparisons use fixed hyperparameters, and that per-dataset tuning reached higher scores on several datasets, with checkpoints promised README.md. Those tuned scores are not in the paper's tables.

Lineage

• Builds on: ImageNet-1K pretraining weights from timm §3.1 app. B; WordNet for the hierarchy §2.2; asymmetric loss for multi-label training §3.2 app. B; knowledge distillation for the KD term §3.4. Papers without a card are named only.
• Compared against: official ImageNet-21K weights for ViT-B-16 and Mixer-B-16 §5.1; ImageNet-1K pretraining §5.2; Open Images pretraining app. J.
• Superseded by: not stated in the paper.