高效获取AI-TOD数据集：自动化下载与校验方案

1次阅读

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

手动下载大型数据集（如 AI-TOD）常面临以下问题：

网络稳定性差：单线程下载在跨国传输时容易因网络波动中断
校验机制缺失：下载后需人工比对哈希值，易出现数据损坏未被发现的情况
效率低下：GB 级数据集通过浏览器下载耗时可能超过 12 小时
恢复成本高：中断后需要重新下载整个文件

requests
优势：API 简洁，社区支持完善
限制：同步阻塞式 IO，需配合线程池实现并发
urllib3
优势：内置连接池管理，底层控制力强
限制：API 较为底层
aiohttp
优势：原生异步支持，高并发场景性能好
限制：需要 asyncio 环境

最终选择：采用 requests+concurrent.futures 组合，平衡开发效率与性能需求

from concurrent.futures import ThreadPoolExecutor

def download_chunk(url, start_byte, end_byte, chunk_file):
    headers = {'Range': f'bytes={start_byte}-{end_byte}'}
    response = requests.get(url, headers=headers, stream=True)
    with open(chunk_file, 'wb') as f:
        for chunk in response.iter_content(1024*8):
            f.write(chunk)

# 使用示例
def parallel_download(url, target_path, threads=4):
    file_size = int(requests.head(url).headers['Content-Length'])
    chunk_size = file_size // threads

    with ThreadPoolExecutor(max_workers=threads) as executor:
        futures = []
        for i in range(threads):
            start = i * chunk_size
            end = start + chunk_size -1 if i < threads-1 else file_size-1
            futures.append(executor.submit(download_chunk, url, start, end, f'{target_path}.part{i}'
            ))
        concurrent.futures.wait(futures)

关键技术点：

检查本地已下载的临时文件
通过 Content-Range 获取剩余字节范围
使用 'a' 模式追加写入文件

def resume_download(url, target_path):
    if os.path.exists(target_path + '.part0'):
        downloaded = os.path.getsize(target_path + '.part0')
        headers = {'Range': f'bytes={downloaded}-'}
    else:
        downloaded = 0
        headers = {}

    response = requests.get(url, headers=headers, stream=True)
    mode = 'ab' if downloaded > 0 else 'wb'

    with open(target_path + '.part0', mode) as f:
        for chunk in response.iter_content(1024*8):
            f.write(chunk)

import hashlib

def verify_file(file_path, expected_hash, algorithm='md5'):
    hash_func = getattr(hashlib, algorithm)()
    with open(file_path, 'rb') as f:
        while chunk := f.read(8192):
            hash_func.update(chunk)
    return hash_func.hexdigest() == expected_hash

from requests.adapters import HTTPAdapter

session = requests.Session()
adapter = HTTPAdapter(
    pool_connections=10,
    pool_maxsize=20,
    max_retries=3
)
session.mount('http://', adapter)
session.mount('https://', adapter)

import time
import math

def download_with_retry(url, max_retries=5):
    for attempt in range(max_retries):
        try:
            response = requests.get(url, timeout=10)
            response.raise_for_status()
            return response
        except Exception as e:
            wait_time = min(10, math.pow(2, attempt))
            time.sleep(wait_time)
    raise Exception(f'Failed after {max_retries} retries')

# 方法 1：全局禁用验证（不推荐）requests.get(url, verify=False)

# 方法 2：指定 CA 证书包
requests.get(url, verify='/path/to/certfile.pem')

关键技巧：

始终使用 stream=True 参数
控制 iter_content 的 chunk_size（推荐 8KB-1MB）
及时关闭响应连接

with requests.get(url, stream=True) as r:
    with open(target_path, 'wb') as f:
        for chunk in r.iter_content(chunk_size=8192):
            if chunk:  # 过滤 keep-alive 数据块
                f.write(chunk)

import os
import hashlib
import concurrent.futures
from typing import Optional

class AITODDownloader:
    def __init__(self, max_workers=4, chunk_size=1024*1024*10):
        self.max_workers = max_workers
        self.chunk_size = chunk_size

    def _download_chunk(self, url, start, end, chunk_file):
        headers = {'Range': f'bytes={start}-{end}'}
        with requests.get(url, headers=headers, stream=True) as r:
            with open(chunk_file, 'wb') as f:
                for chunk in r.iter_content(8192):
                    f.write(chunk)

    def download(self, url, target_path, expected_hash=None):
        # 实现省略，包含完整的多线程下载和校验逻辑
        pass

    @staticmethod
    def verify(file_path, expected_hash, algorithm='md5') -> bool:
        # 实现见前文校验模块
        pass

# 单元测试示例
import unittest
import tempfile

class TestDownloader(unittest.TestCase):
    def test_chunk_download(self):
        with tempfile.NamedTemporaryFile() as tmp:
            downloader = AITODDownloader()
            downloader._download_chunk('http://example.com/file', 0, 1023, tmp.name)
            self.assertGreater(os.path.getsize(tmp.name), 0)

可进一步改造为通用数据集下载框架的方向：