ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

Lance 与 PyTorch 集成实战:用 LanceDataset 构建多模态训练与推理数据流水线

Lance 与 PyTorch 集成实战:用 LanceDataset 构建多模态训练与推理数据流水线 Lance 与 PyTorch 集成实战用 LanceDataset 构建多模态训练与推理数据流水线【免费下载链接】lanceOpen Lakehouse Format for Multimodal AI. Convert from Parquet in 2 lines of code for 100x faster random access, vector index, and data versioning. Compatible with Pandas, DuckDB, Polars, Pyarrow, and PyTorch with more integrations coming..项目地址: https://gitcode.com/GitHub_Trending/la/lance本篇技术指南围绕 Lance 官方 PyTorch 集成文档见 docs/src/integrations/pytorch.md展开讲解如何把 Lance 数据集直接接入 PyTorch 训练/推理循环从一行代码将 HuggingFace 数据集转换为 Lance 数据集到通过lance.torch.data.LanceDataset作为torch.utils.data.IterableDataset使用再到利用lance.sampler在分布式训练中做分片采样。读完本文你将掌握 Lance 数据格式与 PyTorch DataLoader 的完整衔接方式、安全多进程加载方案以及底层实现原理。一、为什么需要在 PyTorch 中使用 Lance 数据集Lance 是一个面向多模态 AI 的开源湖仓格式Open Lakehouse Format for Multimodal AI其核心优势在于列式随机访问、向量索引和数据版本控制。在多模态训练场景中数据集往往包含大量图片、文本、视频等非结构化数据且需要反复、随机地读取子集——这正是 Lance 格式擅长的领域。在 PyTorch 生态中通常的做法是先将数据落盘为图片文件或 NumPy 数组再手写Dataset类读取。而 Lance 提供了一条更直接的路径lance.torch.data.LanceDataset是torch.utils.data.IterableDataset的子类见 python/python/lance/torch/data.py它直接基于 Lance 的扫描器scanner读取数据并自动转换为torch.Tensor训练与推理循环无需关心数据存储细节。二、准备训练数据从 HuggingFace 数据集到 Lance要使用 Lance 训练多模态模型第一步是把数据转换为 Lance 格式。得益于 HuggingFace 集成这一步只需一行代码。import datasets # pip install datasets import lance hf_ds datasets.load_dataset( poloclub/diffusiondb, splittrain, # name2m_first_1k, # 如需更小的子集可取消注释 ) lance.write_dataset(hf_ds, diffusiondb_train.lance)其中lance.write_dataset的data_obj参数直接接受 HuggingFace 数据集见 python/python/lance/dataset.py其 docstring 明确列出可接受类型包括 Huggingface dataset因此无需手动把datasets.Dataset转成 PyArrow Table。该函数还支持以下实用参数完整签名见 python/python/lance/dataset.pymodecreate默认目录已存在则报错、overwrite覆盖为新版本、append追加到最新版本max_rows_per_file默认1024 * 1024单个数据文件的最大行数max_rows_per_group默认1024单个 row group 的最大行数影响随机访问粒度data_storage_version可选stable、2.0、2.1、2.2等新版本编码更高效但需要更新的 Lance 版本才能读取storage_options对象存储如 S3的连接参数。三、基础用法把 Lance 数据集接入 PyTorch 循环数据集转换完成后即可在 PyTorch 的训练/推理循环中使用。文档给出了Unsafe Dataloader基础用法示例import torch import lance.torch.data # 将 Lance 数据集加载为 PyTorch IterableDataset # 仅读取 image 和 prompt 两列。 dataset lance.torch.data.LanceDataset( diffusiondb_train.lance, columns[image, prompt], batch_size128, batch_readahead8, # 控制多线程读取的预读批次数 ) # 创建 PyTorch DataLoader dataloader torch.utils.data.DataLoader(dataset) # 推理循环 for batch in dataloader: inputs, targets batch[prompt], batch[image] outputs model(inputs) ...两点需要注意来自原文档自动类型转换PyTorch 数据集会把数据自动转换为torch.Tensorfork 安全性Lance 不是 fork-safe 的。如果使用多进程请改用 spawn 启动方式get_safe_loader提供的安全 DataLoader 使用 spawn 方法。LanceDataset 的构造参数详解LanceDataset.__init__的完整参数见 python/python/lance/torch/data.py除示例中使用的之外还包括参数默认值说明dataset必填Lance 数据集的 URI、路径或已初始化的lance.LanceDataset对象batch_size必填每次迭代产出的批大小columnsNone要读取的列名列表为None时读取全部列filterNoneSQL 风格过滤表达式如ids 300仅在全量扫描场景下支持samples0若大于 0则只随机采样samples行走maybe_sample路径cacheNone设为True或路径时首次迭代将数据缓存到磁盘后续迭代直接读缓存with_row_idFalse为True时每个 batch 增加一列_rowid表示行 IDbatch_readahead16每个 fragment 在 Rust 线程中预读的批次数控制多线程 IOshard_granularityNone分布式分片粒度fragment或batchto_tensor_fn内置转换自定义 Arrow 批量转 Tensor 的函数samplerNone自定义采样器见第四节默认全量扫描auto_detect_rankTrue自动检测分布式 rank 与 world_size底层类型转换规则从源码看lance.torch.data的默认转换函数_to_tensor见 python/python/lance/torch/data.py支持以下 Arrow 类型到torch.Tensor的映射数值/布尔标量列int、float、bool直接通过to_numpy(zero_copy_onlyFalse)转为 TensorFixedSizeListf16/f32/f64/整数向量列通过_fsl_to_tensor重塑为 2D Tensor形状[batch, dim]见 python/python/lance/torch/data.pyFixedShapeTensor列解包存储后按数值类型转换bfloat16 列Lance 扩展类型lance.bfloat16通过_bf16_to_tensor将底层字节重解释为torch.bfloat16无空值时零拷贝共享 Arrow 缓冲区有空值则克隆后以 NaN 填充见 python/python/lance/torch/data.pyuint64列默认转换为torch.int64uint64_as_int64True。若类型无法转换会抛出ValueError: Only support FixedSizeListf16/bf16/f32/f64 or numeric/bfloat16 values。自定义转换可通过to_tensor_fn参数传入其签名需接受(batch, *, hf_converterNone, use_blob_apiFalse, **kwargs)见 python/python/lance/torch/data.py。测试用例对上述行为有完整验证例如 python/python/tests/torch_tests/test_data.py 覆盖了FixedShapeTensor与FixedSizeList向量列的读取python/python/tests/torch_tests/test_data.py 验证了 bfloat16 列转换为torch.bfloat16。四、Safe Dataloader多进程安全的加载方案Lance 底层是多线程的Rust 侧并行扫描与fork启动的多进程配合不佳。文档明确警告多进程场景下不要使用 fork因为 fork 与多线程不能良好协同。为此lance.torch.data提供了SafeLanceDataset与get_safe_loaderfrom lance.torch.data import SafeLanceDataset, get_safe_loader dataset SafeLanceDataset(temp_lance_dataset) # 使用 spawn 方法以避免 fork-safe 问题 loader get_safe_loader( dataset, num_workers2, batch_size16, drop_lastFalse, ) total_samples 0 for batch in loader: total_samples batch[id].shape[0]其实现要点见 python/python/lance/torch/data.pySafeLanceDataset继承torch.utils.data.Dataset而非 IterableDataset在主进程中仅做轻量元数据预载_safe_preload只统计行数每个 worker 进程首次取数时才延迟打开数据集__getitems__中的_ds is None分支从而避免在 fork 前初始化复杂的 Rust 扫描器get_safe_loader强制使用torch.multiprocessing.get_context(spawn)作为multiprocessing_context并将persistent_workers默认设为True同时保留用户传入参数的最高优先级SafeLanceDataset也接受dataset_options如version、storage_options且 worker 进程重开数据集时会保留这些选项——python/python/tests/torch_tests/test_data.py 中的回归测试专门验证了这一点。五、分布式训练与 lance.sampler 组合控制采样策略lance.torch.data.LanceDataset可以与lance.sampler.Sampler类组合来控制采样策略。例如在分布式训练环境中使用lance.sampler.ShardedFragmentSampler将不同 fragment 分配给不同进程不指定 sampler 时默认为全量扫描。from lance.sampler import ShardedFragmentSampler from lance.torch.data import LanceDataset # 将 Lance 数据集加载为 PyTorch IterableDataset # 仅读取 image 和 prompt 两列。 dataset LanceDataset( diffusiondb_train.lance, columns[image, prompt], batch_size128, batch_readahead8, # 控制多线程读取的预读批次数 samplerShardedFragmentSampler( rank1, # 当前进程的 rank world_size8, # 总进程数 ), )文档列出的可用 sampler 有lance.sampler.ShardedFragmentSamplerlance.sampler.ShardedBatchSamplerSampler 家族与实现原理lance.sampler模块见 python/python/lance/sampler.py围绕抽象基类Sampler其__call__方法产出一个pyarrow.RecordBatch生成器见 python/python/lance/sampler.py提供了多种实现FullScanSampler默认采样器按顺序扫描整个数据集python/python/lance/sampler.pyShardedFragmentSampler按rank与world_size对 fragment 分片即每个进程处理ds.fragments[rank::world_size]的子集。其 docstring 指出数据集较大时它比ShardedBatchSampler更高效python/python/lance/sampler.py。它还支持randomize与seed参数打乱 fragment 顺序并可调用set_epoch(epoch)实现每轮 epoch 的 shuffleShardedBatchSampler按 batch 粒度分片每个进程取第 N 个 batchN 为 world_size。无过滤条件时各进程只加载分配给自己的行区间有过滤条件时每个进程都会扫描完整过滤结果再丢弃不属于自己的行并用基于堆的近似蓄水池采样reservoir samplingk256打乱顺序python/python/lance/sampler.pyShardedFixedBatchSampler面向已知总行数的静态数据集为每个 rank 分配连续的、不重叠的索引区间见 python/python/lance/sampler.py。自动检测分布式环境ShardedFragmentSampler与ShardedBatchSampler都提供了from_torch(randomize, seed)类方法直接从torch.distributed.get_rank()/get_world_size()自动推断 rank 与 world_size见 python/python/lance/sampler.py 和 python/python/lance/sampler.py。即使不显式传 samplerLanceDataset在auto_detect_rankTrue默认时也会通过lance.torch.dist模块自动检测全局 rank 与世界大小get_global_rank与get_global_world_size综合了torch.distributed的进程级信息和 DataLoader worker 级信息get_worker_info()以rank dist_rank * mp_world_size mp_rank的方式计算全局 rank见 python/python/lance/torch/dist.py。因此在多进程 多 worker 的复杂场景下也能正确分片。测试用例 python/python/tests/torch_tests/test_data.py 验证了两种 sampler 的分片行为例如 rank1、world_size2 时ShardedFragmentSampler只取编号为奇数的文件片段python/python/tests/torch_tests/test_data.py 验证了带过滤条件的ShardedBatchSampler分片与随机化逻辑。六、进阶用法与实用技巧1. 随机采样训练子集samples当samples大于 0 时LanceDataset会走lance.sampler.maybe_sample路径见 python/python/lance/torch/data.py若采样数不小于数据总量则全量扫描否则结合随机take随机读与分块采样实现既保证随机性又控制随机读次数max_takes默认 2048见 python/python/lance/sampler.py。这在需要固定规模子集做调试或对比实验时非常实用。2. 过滤与 row_idfilter参数支持 SQL 风格表达式如ids 300仅在无采样、无 batch 粒度分片时可用python/python/lance/torch/data.py。with_row_idTrue时每个 batch 会附带_rowid列可配合 Lance 的随机访问能力。3. 磁盘缓存cache设cacheTrue时首次迭代会把数据流写入磁盘缓存CachedDataset后续迭代直接从缓存读取避免重复扫描原始数据python/python/lance/torch/data.py。测试 python/python/tests/torch_tests/test_data.py 验证了删除原始数据后仍可从缓存迭代。4. Blob 列与自定义转换当数据列以 Large Blob 编码存储schema 元数据lance-encoding:blobtrue时LanceDataset会自动检测 blob 列并通过take_blobs惰性读取同时强制开启with_row_id见 python/python/lance/torch/data.py 与 python/python/lance/torch/data.py。此时需要用户通过to_tensor_fn自行把 blob 内容转换为 Tensorpython/python/tests/torch_tests/test_data.py 给出了完整的自定义转换示例。七、使用注意事项进程模型Lance 内部是多线程的fork 会继承这些线程并可能导致死锁或数据损坏。多进程场景应使用 spawnget_safe_loader已默认强制 spawn类型覆盖范围内置转换仅支持数值/布尔标量、FixedSizeList向量float/integer、FixedShapeTensor与 bfloat16字符串、嵌套结构等其他类型需自定义to_tensor_fn采样与过滤的兼容性filter与samples不能同时使用会抛出NotImplementedError测试见 python/python/tests/torch_tests/test_data.pyShardedBatchSampler不支持with_row_idrank/world_size参数已弃用直接传构造参数会触发DeprecationWarningpython/python/lance/torch/data.py推荐使用sampler参数或依赖auto_detect_rank自动检测。八、更多相关资源PyTorch 集成原始文档docs/src/integrations/pytorch.md集成总览含 DataFusion、DuckDB、TensorFlow、Ray 等生态docs/src/integrations/index.md核心实现python/python/lance/torch/data.py采样器实现python/python/lance/sampler.py分布式工具python/python/lance/torch/dist.py设备选择工具preferred_device自动选择 CUDA/MPS/CPUpython/python/lance/torch/init.py测试用例python/python/tests/torch_tests/test_data.py【免费下载链接】lanceOpen Lakehouse Format for Multimodal AI. Convert from Parquet in 2 lines of code for 100x faster random access, vector index, and data versioning. Compatible with Pandas, DuckDB, Polars, Pyarrow, and PyTorch with more integrations coming..项目地址: https://gitcode.com/GitHub_Trending/la/lance创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表