本文介绍使用 os.scandir() 替代 os.listdir() 高效遍历大目录并随机抽取 k 个文件的方法，避免内存暴涨与编辑器（如 VS Code）崩溃，特别适用于构建图像 DataLoader 的场景。
本文介绍使用 `os.scandir()` 替代 `os.listdir()` 高效遍历大目录并随机抽取 k 个文件的方法，避免内存暴涨与编辑器（如 vs code）崩溃，特别适用于构建图像 dataloader 的场景。
在处理大规模图像数据集（例如数万张 PNG 文件）时，直接调用 os.listdir(path) 会一次性将所有文件名加载到内存中，不仅耗时，还极易引发内存压力——这正是 VS Code 崩溃的常见原因。而 os.scandir() 是 Python 3.5+ 引入的底层优化接口：它返回轻量级 DirEntry 对象流式迭代器，仅在访问 .name 或 .is_file() 等属性时才解析必要元数据，显著降低内存开销和启动延迟。
以下是一个生产就绪的实现方案，用于从 data/images/ 目录中高效、随机、无重复地选取 64 张图像：
import os
import random
def sample_k_files_from_dir(directory: str, k: int = 64, extension: str = ".png") -> list[str]:
"""
从指定目录中随机采样 k 个符合扩展名的文件（不预加载全部文件名）
返回相对路径列表，便于后续构造完整路径或集成至 DataLoader
"""
# 使用 scandir 流式获取文件条目，过滤出目标类型文件
entries = os.scandir(directory)
filenames = [
entry.name
for entry in entries
if entry.is_file() and entry.name.lower().endswith(extension.lower())
]
# 安全校验：确保有足够文件可选
if len(filenames) < k:
raise ValueError(f"Directory '{directory}' contains only {len(filenames)} {extension} files, less than requested k={k}")
# 随机无放回抽样（时间复杂度 O(k)，优于 shuffle 全列表）
selected_names = random.sample(filenames, k=k)
return selected_names
# 使用示例
output_dir = "data/images"
sampled_files = sample_k_files_from_dir(output_dir, k=64)
# 构造完整路径并加载（可对接 torch.utils.data.Dataset 或 PIL）
for fname in sampled_files:
img_path = os.path.join(output_dir, fname)
# 示例：用 PIL 打开图像
# from PIL import Image; img = Image.open(img_path)✅ 关键优势说明：
- ✅ 内存友好：os.scandir() 不缓存整个目录结构，filenames 列表仅存储文件名字符串（而非完整路径或 inode 信息），大幅压缩内存占用；
- ✅ 类型安全：通过 entry.is_file() 显式过滤，自动跳过子目录、符号链接及系统隐藏文件；
- ✅ 扩展灵活：支持按后缀（如 .png, .jpg）精确筛选，避免误选非图像文件；
- ✅ 健壮性强：内置样本数量校验，防止 random.sample() 因 k > n 抛出 ValueError。
⚠️ 注意事项：
- 若目录中文件数量极大（如 >10⁶），即使使用 scandir，构建 filenames 列表仍需 O(n) 时间。此时可考虑蓄水池采样（Reservoir Sampling） 实现 O(n) 时间 + O(k) 空间的真正流式随机抽样（无需先收集全部文件名）。如需该进阶方案，可进一步封装为生成器版本；
- os.scandir() 在使用后建议显式关闭迭代器以释放系统资源（尤其在循环或长期运行服务中），但本例中列表推导式会自动耗尽迭代器，通常无需手动 entries.close()；
- Windows 下 scandir 性能提升尤为显著；Linux/macOS 也有稳定收益，但差异略小。
综上，os.scandir() + 列表推导 + random.sample() 的组合，是在保持代码简洁性与工程可读性的前提下，兼顾性能、可靠性和可维护性的最优实践。它完美契合深度学习训练中“按需加载小批量图像”的典型需求，是构建高效 DataLoader 的坚实基础。
