高效获取bdd100k数据集图像分割数据的完整解决方案

1次阅读

共计 2277 个字符，预计需要花费 6 分钟才能阅读完成。

BDD100K 是伯克利大学发布的自动驾驶领域大规模数据集，包含 10 万张图像及其标注信息。其中图像分割数据包含像素级的道路场景标注，对自动驾驶算法训练至关重要。然而，下载这个超过 100GB 的数据集时，开发者常遇到以下痛点：

网络连接不稳定导致下载中断
单线程下载速度慢（平均仅 2 -3MB/s）
重试成本高，失败后需从头开始下载
本地存储空间不足时难以分批次下载

我们采用三种核心技术解决上述问题：

分片下载：将大文件拆分为多个小片段（如 100MB/ 片），独立下载后合并
断点续传：记录已完成片段，中断后自动跳过已下载部分
并行处理：通过线程池并发下载多个片段，充分利用带宽

import os
import requests
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

def download_file(url: str, save_path: Path, chunk_size=1024*1024, max_workers=8):
    """
    分片并行下载文件
    :param url: 文件 URL
    :param save_path: 本地保存路径
    :param chunk_size: 分片大小（字节）:param max_workers: 最大并发数
    """
    # 获取文件总大小
    headers = requests.head(url).headers
    total_size = int(headers.get('content-length', 0))

    # 计算分片数量
    chunks = range(0, total_size, chunk_size)

    # 创建临时目录
    temp_dir = save_path.parent / f"{save_path.name}.temp"
    temp_dir.mkdir(exist_ok=True)

    def download_chunk(chunk_index):
        start = chunk_index * chunk_size
        end = min(start + chunk_size - 1, total_size - 1)

        # 设置 Range 头实现分片下载
        headers = {'Range': f'bytes={start}-{end}'}
        chunk_file = temp_dir / f"chunk_{chunk_index:04d}"

        if chunk_file.exists():
            return chunk_index

        try:
            with requests.get(url, headers=headers, stream=True) as r:
                r.raise_for_status()
                with open(chunk_file, 'wb') as f:
                    for chunk in r.iter_content(8192):
                        f.write(chunk)
            return chunk_index
        except Exception as e:
            if chunk_file.exists():
                chunk_file.unlink()
            raise e

    # 并行下载所有分片
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = [executor.submit(download_chunk, i) for i in range(len(chunks))]

        for future in as_completed(futures):
            try:
                chunk_index = future.result()
                print(f"Downloaded chunk {chunk_index+1}/{len(chunks)}")
            except Exception as e:
                print(f"Error downloading chunk: {e}")

    # 合并分片
    with open(save_path, 'wb') as out_file:
        for i in range(len(chunks)):
            chunk_file = temp_dir / f"chunk_{i:04d}"
            with open(chunk_file, 'rb') as in_file:
                out_file.write(in_file.read())
            chunk_file.unlink()

    temp_dir.rmdir()

我们在 AWS EC2（100Mbps 带宽）上测试不同并发数的表现：