Fork me on github

从ImageNet-1k制作ImageNet-Subset

# -*- coding: utf-8 -*-
"""
Create a subset of the ImageNet-1k dataset.
"""
import os
import numpy as np
import torchvision.datasets as datasets
SUBSET_SIZE = 100
seed = 1993
SOURCE_PATH = "/dataset/ImageNet-1k"
TARGET_SIZE = f"/dataset/ImageNet-{SUBSET_SIZE}_seed{seed}"
# 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_SIZE, "train")
output_val_dir = os.path.join(TARGET_SIZE, "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!")
posted @   fang-d  阅读(45)  评论(0编辑  收藏  举报
相关博文:
阅读排行:
· 无需6万激活码!GitHub神秘组织3小时极速复刻Manus,手把手教你使用OpenManus搭建本
· C#/.NET/.NET Core优秀项目和框架2025年2月简报
· Manus爆火,是硬核还是营销?
· 终于写完轮子一部分:tcp代理 了,记录一下
· 【杭电多校比赛记录】2025“钉耙编程”中国大学生算法设计春季联赛(1)
点击右上角即可分享
微信分享提示