使用 Lance 格式管理 PyTorch 深度学习工件模型权重的保存、版本化与加载【免费下载链接】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 开源仓库中的官方示例 docs/src/examples/python/artifact_management.md 为蓝本演示如何把 PyTorch 模型的权重state dict保存为 Lance 数据集利用 Lance 的零拷贝Zero-copy自动版本化能力在同一份数据集中管理多个模型检查点并支持按版本号回读、还原为原始形状的权重张量。读完本文你将掌握一套可以替代每个检查点一个文件夹的深度学习工件管理方案并理解其背后的数据模型与 API 实现。核心思路用三列表模拟 PyTorch 的 state dict在 PyTorch 中模型的权重通过state_dict()获取它本质上是一个OrderedDict键是权重名称如conv1.weight值是对应的权重张量torch.Tensor。Lance 是列式存储格式没有直接的字典结构因此示例文档采用了一个巧妙的三列表模型来等价表达 state dict列名类型含义namepa.string()权重的名称对应 state dict 的键valuepa.list_(pa.float64(), -1)展平flatten后的权重数值列表shapepa.list_(pa.int64(), -1)权重的原始形状用于加载时重塑reshape回张量其中value和shape都使用变长列表类型-1表示不限定长度因为不同层权重的维度数不同例如 4 维卷积核与 1 维偏置。这样每一行就对应 state dict 中的一个键值对整个数据集就是一个可版本化的权重仓库。环境准备与 Schema 定义在开始之前需要先安装 Lance、PyArrow 与 PyTorchpip install lance pyarrow torch然后导入所需模块并定义全局 Schemaimport os import shutil import lance import pyarrow as pa import torch from collections import OrderedDict GLOBAL_SCHEMA pa.schema( [ pa.field(name, pa.string()), pa.field(value, pa.list_(pa.float64(), -1)), pa.field(shape, pa.list_(pa.int64(), -1)), # Is a list with variable shape because weights can have any number of dims ] )这份GLOBAL_SCHEMA决定了权重表的结构name列存权重名、value列存展平后的权重值、shape列存原始形状在后续写入lance.write_dataset和读取lance.dataset时保持列式结构的统一。保存模型将 state dict 写入 Lance 数据集把每个权重转成一行 RecordBatch首先编写一个生成器函数遍历 state dict 中的每个参数将其展平并连同名称、形状一起封装成pyarrow.RecordBatchdef _save_model_writer(state_dict): Yields a RecordBatch for each parameter in the model state dict for param_name, param in state_dict.items(): param_shape list(param.size()) param_value param.flatten().tolist() yield pa.RecordBatch.from_arrays( [ pa.array( [param_name], pa.string(), ), pa.array( [param_value], pa.list_(pa.float64(), -1), ), pa.array( [param_shape], pa.list_(pa.int64(), -1), ), ], [name, value, shape], )这里的关键点是param.flatten().tolist()将任意维度的张量展平为一维列表与list(param.size())记录原始形状两者在加载阶段互为逆操作。利用 overwrite 模式实现版本化写入接下来是核心的save_model函数。它负责三种场景目标数据集不存在时直接创建已存在且开启版本化时追加为新版本已存在且不开启版本化时删除后重写def save_model(state_dict: OrderedDict, file_name: str, versionFalse): Saves a PyTorch model in lance file format Args: state_dict (OrderedDict): Model state dict file_name (str): Lance model name version (bool): Whether to save as a new version or overwrite the existing versions, if the lance file already exists # Create a reader reader pa.RecordBatchReader.from_batches( GLOBAL_SCHEMA, _save_model_writer(state_dict) ) if os.path.exists(file_name): if version: # If we want versioning, we use the overwrite mode to create a new version lance.write_dataset( reader, file_name, schemaGLOBAL_SCHEMA, modeoverwrite ) else: # If we dont want versioning, we delete the existing file and write a new one shutil.rmtree(file_name) lance.write_dataset(reader, file_name, schemaGLOBAL_SCHEMA) else: # If the file doesnt exist, we write a new one lance.write_dataset(reader, file_name, schemaGLOBAL_SCHEMA)从仓库源码 python/python/lance/dataset.py 可以看到lance.write_dataset的mode参数支持三种取值create默认新建数据集若路径已存在则报错overwrite创建新的快照版本——这正是示例文档中版本化的实现方式。每次 overwrite 并不会物理删除旧数据而是提交一个新的 manifest 版本旧版本仍可通过版本号访问append创建新版本内容是输入数据与最新版本的拼接。此外write_dataset接受 Pandas DataFrame、PyArrow Table、Dataset、Scanner 或RecordBatchReader等 Reader-like 数据源。示例中使用pa.RecordBatchReader.from_batches(GLOBAL_SCHEMA, _save_model_writer(state_dict))流式生成批量数据对于大模型可以避免一次性把所有权重堆进内存。write_dataset还提供了一些对大型权重数据集很有用的参数例如max_rows_per_group默认 1024控制单个文件内 group 的大小与max_bytes_per_file默认 90 GiB对象存储单文件软上限可按需调整以优化随机读取性能。版本化带来了什么假设你在新数据上微调了模型但又不想覆盖旧检查点传统做法是新建一个文件夹存放新权重而使用 Lance 后只需再次以modeoverwrite写入权重数据集会自动产生一个新版本。这样所有检查点收敛到同一个数据集目录不再散落多个文件夹由于 Lance 的零拷贝自动版本化未变更的数据文件在版本间被共享复用空间开销极小随时可以通过版本号回溯到任意一次训练结果。加载模型从 Lance 数据集还原权重加载是保存的逆过程核心在于利用shape列把展平的数值重塑回原始形状。示例文档将其拆分为三个层次分明的函数。单条权重还原def _load_weight(weight: dict) - torch.Tensor: Converts a weight dict to a torch tensor return torch.tensor(weight[value], dtypetorch.float64).reshape(weight[shape])从 Lance 数据集取回的一行数据在to_pylist()后是{name: ..., value: [...], shape: [...]}形式的字典_load_weight负责把value列表转成torch.Tensor并按shape重塑。你还可以在此处扩展一个dtype参数以支持float16、bfloat16等量化或混合精度场景避免加载时默认使用float64。组装 state dictdef _load_state_dict(file_name: str, version: int 1, map_locationNone) - OrderedDict: Reads the model weights from lance file and returns a model state dict If the model weights are too large, this function will fail with a memory error. Args: file_name (str): Lance model name version (int): Version of the model to load map_location (str): Device to load the model on Returns: OrderedDict: Model state dict ds lance.dataset(file_name, versionversion) weights ds.take([x for x in range(ds.count_rows())]).to_pylist() state_dict OrderedDict() for weight in weights: state_dict[weight[name]] _load_weight(weight).to(map_location) return state_dict这里的lance.dataset(file_name, versionversion)是版本化回读的关键。查看仓库实现 python/python/lance/init.pyversion参数既可以是整数版本号也可以是字符串标签tag不传时默认读取最新版本。也就是说你可以把一次训练的权重保存为version1另一次保存为version2之后用lance.dataset(path, version2)精确取回第二次的结果。在取数阶段ds.take([x for x in range(ds.count_rows())])一次性取回全部行count_rows()返回数据集行数即权重总数take按行索引随机读取并返回pa.Table最后to_pylist()转为字典列表。对于常见的 ResNet 规模模型数百个参数项、每项数万浮点数这套流程足够高效take的随机访问能力也来自 Lance 列式格式的索引设计这正是项目描述中100x faster random access所指的能力之一。需要说明的是示例文档明确指出该函数假设权重总量可以放入内存因此不涉及分片sharding。若权重超过单机内存需要将take改为按批次扫描例如使用scanner流式读取。加载进模型def load_model( model: torch.nn.Module, file_name: str, version: int 1, map_locationNone ): Loads the model weights from lance file and sets them to the model Args: model (torch.nn.Module): PyTorch model file_name (str): Lance model name version (int): Version of the model to load map_location (str): Device to load the model on state_dict _load_state_dict(file_name, versionversion, map_locationmap_location) model.load_state_dict(state_dict)最外层只需传入模型实例、数据集路径、版本号与设备位置内部自动完成取数据 → 重塑 → 组装 state dict → 灌入模型的全链路。map_location参数透传给torch.Tensor.to()可用于把权重加载到 GPU如cuda:0或指定设备。完整使用示例ResNet 权重的保存、版本化与校验将上述函数串起来即可完整复现示例文档描述的工作流——加载一个预训练 ResNet保存其权重修改后作为新版本保存再按版本回读并校验权重一致性import torchvision.models as models # 1. 加载预训练 ResNet 并获取 state dict model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) state_dict model.state_dict() # 2. 保存为 Lance 权重数据集版本 1 save_model(state_dict, resnet18.lance) # 3. 模拟一次微调修改部分权重后保存为新版本版本 2 new_state_dict OrderedDict(state_dict) new_state_dict[fc.weight] new_state_dict[fc.weight] * 0.5 save_model(new_state_dict, resnet18.lance, versionTrue) # 4. 回读版本 1 的权重加载到新模型并校验 model_v1 models.resnet18() load_model(model_v1, resnet18.lance, version1) loaded_state_dict model_v1.state_dict() for key in state_dict.keys(): assert torch.equal(state_dict[key], loaded_state_dict[key]), fMismatch: {key} print(Weights restored from version 1 match the original model exactly.)校验通过即证明展平 → 列式存储 → 重塑这条链路对权重的数值是无损的这也是示例文档verifying if the weights are still indeed the same的核心验证目标。进一步探索版本查询与按版本回看Lance Python SDK 还提供了与版本化配套的查询能力可帮助你管理权重历史LanceDataset.versions()python/python/lance/dataset.py返回数据集的全部版本列表每个版本包含版本号与提交时间戳LanceDataset.latest_versionpython/python/lance/dataset.py返回最新版本号LanceDataset.checkout_version(version)python/python/lance/dataset.py复用当前缓存切换到指定版本适合在同一进程中连续查看多个检查点。仓库测试用例 python/python/tests/test_dataset.py 中大量使用了checkout_version与版本迭代逻辑可以当作版本化 API 的权威使用参考。更完整的版本管理与分支能力包括标签、asof时间点读取可参考 docs/src/guide/read_and_write.md 与 docs/src/guide/versioning.md。总结与适用边界整个方案对外只暴露两个函数save_model与load_model。前者把 PyTorch 的OrderedDict编码成名称 展平值 形状的三列 Lance 数据集并通过modeoverwrite实现同一路径下的多版本管理后者通过lance.dataset(..., versionN)take按版本回读、按shape重塑最终调用model.load_state_dict完成加载。需要强调的边界条件示例文档中均已注明该方案目前属于实验性方法适合在权重可完整装入内存的规模下使用超大模型需自行扩展为分片/流式读取权重类型以 PyTorch 张量为前提value列固定使用float64如需float16等精度需扩展_load_weight的dtype参数版本化依赖 Lance 的overwrite语义与零拷贝版本机制因此无需为每个检查点维护独立目录这也是相比传统权重保存方式更精简streamlined的核心价值所在。【免费下载链接】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),仅供参考