从ImageNet-1k制作ImageNet-Subset

# -*- coding: utf-8 -*-
"""
Create a subset of the ImageNet-1k dataset.

Ref: Hou, Saihui, et al. "Learning a Unified Classifier Incrementally via Rebalancing." 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR). IEEE, 2019.
URL: https://github.com/hshustc/CVPR19_Incremental_Learning/blob/master/imagenet-class-incremental/gen_imagenet_subset.py
"""
import os
import numpy as np
from torchvision import datasets

SUBSET_SIZE = 100
SEED = 1993
SOURCE_PATH = "/dataset/ImageNet-1k"
TARGET_PATH = f"/dataset/ImageNet-{SUBSET_SIZE}_seed{SEED}"

SOURCE_PATH = os.path.expanduser(SOURCE_PATH)
TARGET_PATH = os.path.expanduser(TARGET_PATH)

# Data loading code
train_dataset = datasets.ImageNet(SOURCE_PATH, split="train")
classes = train_dataset.wnids

# Randomly select a subset of classes
np.random.seed(SEED)
subset_classes = np.random.choice(classes, SUBSET_SIZE, replace=False)
subset_classes.sort()
print("the number of subset classes: {}".format(len(subset_classes)))
print(subset_classes)

# Create links to the subset folder
source_train_dir = os.path.join(SOURCE_PATH, "train")
source_val_dir = os.path.join(SOURCE_PATH, "val")
output_train_dir = os.path.join(TARGET_PATH, "train")
output_val_dir = os.path.join(TARGET_PATH, "val")
os.makedirs(output_train_dir, exist_ok=True)
os.makedirs(output_val_dir, exist_ok=True)

for cls in subset_classes:
   os.symlink(os.path.join(source_train_dir, cls), os.path.join(output_train_dir, cls))
   os.symlink(os.path.join(source_val_dir, cls), os.path.join(output_val_dir, cls))

print("Done!")