如何用Python os模块不加载所有文件名随机选k个文件？解决VSCode崩溃问题

直接用os.listdir()会一次性把所有文件名加载到内存，文件量一大就容易导致崩溃。但可以通过os.scandir()（迭代式遍历）结合蓄水池抽样算法实现不加载全部文件名的随机选取，或者搭配glob.iglob()来做同样的事，以下是具体实现：

方法1：os.scandir() + 蓄水池抽样 #

os.scandir()是迭代器模式，不会一次性返回所有文件条目，内存占用极低。遍历过程中用蓄水池抽样保证随机选取64个PNG文件，不用先存储所有文件名：

import os
import random

def random_sample_pngs(dir_path, sample_size=64):
    sample = []
    count = 0
    # 迭代遍历目录条目
    with os.scandir(dir_path) as entries:
        for entry in entries:
            # 筛选PNG文件
            if entry.is_file() and entry.name.lower().endswith('.png'):
                count += 1
                # 蓄水池抽样逻辑：前64个直接加入，之后随机替换样本池中的元素
                if len(sample) < sample_size:
                    sample.append(entry.path)
                else:
                    r = random.randint(0, count - 1)
                    if r < sample_size:
                        sample[r] = entry.path
    return sample

方法2：glob.iglob() + 蓄水池抽样 #

glob.iglob()同样是迭代式返回匹配的文件路径，适合直接筛选PNG格式，用法和上面类似：

import glob
import os
import random

def random_sample_pngs_glob(dir_path, sample_size=64):
    sample = []
    count = 0
    # 迭代获取所有PNG文件路径
    for file_path in glob.iglob(os.path.join(dir_path, '*.png'), recursive=False):
        count += 1
        if len(sample) < sample_size:
            sample.append(file_path)
        else:
            r = random.randint(0, count - 1)
            if r < sample_size:
                sample[r] = file_path
    return sample

如果数据集有子目录，把recursive参数改成True即可。

适配DataLoader的实现 #

拿到随机样本路径后，自定义Dataset类再包装成DataLoader即可：

from torch.utils.data import Dataset, DataLoader
from PIL import Image
import torchvision.transforms as transforms

class RandomPNGDataset(Dataset):
    def __init__(self, file_paths, transform=None):
        self.file_paths = file_paths
        self.transform = transform or transforms.ToTensor()
    
    def __len__(self):
        return len(self.file_paths)
    
    def __getitem__(self, idx):
        img = Image.open(self.file_paths[idx]).convert('RGB')
        return self.transform(img)

# 获取随机样本
sample_paths = random_sample_pngs('./your_dataset_dir')
# 创建DataLoader
dataloader = DataLoader(RandomPNGDataset(sample_paths), batch_size=64, shuffle=True)

每次训练前调用random_sample_pngs就能获取新的随机批次样本。

内容的提问来源于stack exchange，提问作者Arpit Anand