Awesome

Soft Label Pruning for Large-scale Dataset Distillation (<ins>LPLD</ins>)

[Paper | BibTex | Google Drive]

Official Implementation for "Are Large-scale Soft Labels Necessary for Large-scale Dataset Distillation?", published at NeurIPS'24.

Lingao Xiao,  Yang He

Abstract: In ImageNet-condensation, the storage for auxiliary soft labels exceeds that of the condensed dataset by over 30 times. However, are large-scale soft labels necessary for large-scale dataset distillation? In this paper, we first discover that the high within-class similarity in condensed datasets necessitates the use of large-scale soft labels. This high within-class similarity can be attributed to the fact that previous methods use samples from different classes to construct a single batch for batch normalization (BN) matching. To reduce the within-class similarity, we introduce class-wise supervision during the image synthesizing process by batching the samples within classes, instead of across classes. As a result, we can increase within-class diversity and reduce the size of required soft labels. A key benefit of improved image diversity is that soft label compression can be achieved through simple random pruning, eliminating the need for complex rule-based strategies. Experiments validate our discoveries. For example, when condensing ImageNet-1K to 200 images per class, our approach compresses the required soft labels from 113 GB to 2.8 GB (40x compression) with a 2.6% performance gain.

Installation

Donwload repo:

git clone https://github.com/he-y/soft-label-pruning-for-dataset-distillation.git LPLD
cd LPLD

Create pytorch environment:

conda env create -f environment.yml
conda activate lpld

Download all datasets and labels

Method 1: Automatic Downloading

# sh download.sh [true|false]
sh download.sh false

• true|false meaning whether to download only 40x compressed labels or all labels. (default: false, download all labels)

Method 2: Manual Downloading

Download manually from Google Drive, and place downloaded files in the following structure:

.
├── README.md
├── recover
│   └── model_with_class_bn
│       └── [put Models-with-Class-BN here]
│   └── validate_result
│       └── [put Distilled-Datast here]
├── relabel_and_validate
│   └── syn_label_LPLD
│       └── [put Labels here]

You will find following after downloading

Model with Class BN

Distilled Image Dataset

Previous Soft Labels vs Ours

• full labels for ImageNet-21K are too large to upload; nevertheless, we provide the 40x pruned labels.
• labels for other compression ratios are provided in google drive, or refer README: Usage to generate the labels.

Necessary Modification for Pytorch

Modify PyTorch source code torch.utils.data._utils.fetch._MapDatasetFetcher to support multi-processing loading of soft label data and mix configurations.

class _MapDatasetFetcher(_BaseDatasetFetcher):
    def fetch(self, possibly_batched_index):
        if hasattr(self.dataset, "mode") and self.dataset.mode == 'fkd_load':
            if hasattr(self.dataset, "G_VBSM") and self.dataset.G_VBSM:
                pass # G_VBSM: uses self-decoding in the training script
            elif hasattr(self.dataset, "use_batch") and self.dataset.use_batch:
                mix_index, mix_lam, mix_bbox, soft_label = self.dataset.load_batch_config_by_batch_idx(possibly_batched_index[0])
            else:
                mix_index, mix_lam, mix_bbox, soft_label = self.dataset.load_batch_config(possibly_batched_index[0])

        if self.auto_collation:
            if hasattr(self.dataset, "__getitems__") and self.dataset.__getitems__:
                data = self.dataset.__getitems__(possibly_batched_index)
            else:
                data = [self.dataset[idx] for idx in possibly_batched_index]
        else:
            data = self.dataset[possibly_batched_index]

        if hasattr(self.dataset, "mode") and self.dataset.mode == 'fkd_load':
            # NOTE: mix_index, mix_lam, mix_bbox can be None
            mix_index_cpu = mix_index.cpu() if mix_index is not None else None
            return self.collate_fn(data), mix_index_cpu, mix_lam, mix_bbox, soft_label.cpu()
        else:
            return self.collate_fn(data)

Reproduce Results for 40x compression ratio

To reproduce the [Table] for 40x compression ratio, run the following code:

cd relabel_and_validate
bash scripts/reproduce/main_table_in1k.sh
bash scripts/reproduce/main_table_tiny.sh
bash scripts/reproduce/main_table_in21k.sh

NOTE: validation directory (val_dir) in config files (relabel_and_validate/cfg/reproduce/CONFIG_FILE) should be changed to correct path on your device.

To Reproduce Results for other compression ratios

Please refer to README: Usage for details, including three modules.

Table Results (Google Drive)

Related Repos

Our code is mainly related to the following papers and repos:

• Squeeze, Recover and Relabel: Dataset Condensation at ImageNet Scale From A New Perspective
• ImageNet-21K Pretraining for the Masses

Citation

@inproceedings{xiao2024lpld,
  title={Are Large-scale Soft Labels Necessary for Large-scale Dataset Distillation?},
  author={Lingao Xiao and Yang He},
  booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems},
  year={2024}
}