coremltools Palettization 技术全解析:权重聚类查找表压缩的原理、粒度控制与三套算法实战
发布时间:2026/10/13 9:28:03 锦皓数字建站

AI 应用机器学习模型优化【免费下载链接】coremltoolsCore ML tools contain supporting tools for Core ML model conversion, editing, and validation.项目地址https://gitcode.com/gh_mirrors/co/coremltools点击查看免费下载Palettization也称为 weight clustering / 权重重聚类是 Core ML 模型压缩的核心手段之一它将模型的 float 权重聚类为少量簇心centroid构建查找表Lookup TableLUT并用指向簇心的索引代替原始权重从而显著减小模型体积。本文以 docs-guides/source/opt-palettization-overview.md 为骨架结合coremltools.optimize的 Core ML 与 PyTorch 两侧源码实现系统讲解 palettization 的数学原理、粒度配置、向量聚类、LUT 量化以及 K-Means / Sensitive K-Means / Differentiable K-Means 三套算法的选型与完整实战代码。读完本文你将掌握在 coremltools 中为 mlprogram 模型或 PyTorch 模型配置并执行 1/2/3/4/6/8 位权重聚类压缩的完整方案。Palettization 是什么用簇心索引替代 float 权重Palettization 通过聚类把模型的 float 权重压缩成「查找表LUT 索引」两部分数值相近的权重被归入同一个簇统一用该簇的簇心centroid表示。原始权重矩阵因此被转换为一张索引表表中每个元素指向其所属簇在 LUT 中的位置。以一个简单例子说明出自 palettize_weights 函数文档float 权重向量{0.3, 0.3, 0.5, 0.5}可以用 1-bit LUT{0.3, 0.5}压缩权重向量被替换为 1-bit 索引{0, 0, 1, 1}。运行时Core ML 依据 LUT 与索引现场重建 float 权重再参与计算权重精度由 float 32/16 降低为非均匀的簇心值——这正是「非均匀精度降低non-uniform lowering of precision」的含义。支持的位宽N {1, 2, 3, 4, 6, 8}N为 palettization 使用的比特数每个 LUT 拥有2^N个条目。coremltools 支持N{1,2,3,4,6,8}这一集合在源码中作为硬约束存在Core ML 侧 _quantization_passes.py 定义了_SUPPORTED_NBITS (1, 2, 3, 4, 6, 8)PyTorch 侧 K-Means 实现同样遵循该集合。当传入集合之外的nbits时会抛出Invalid nbits. Support (1, 2, 3, 4, 6, 8)类错误。粒度Granularity从 per_tensor 到 per_grouped_channel上图展示的是per_tensor粒度整张权重张量共享一个LUT。对于大矩阵单一 LUT 会带来较高的近似误差。为此从iOS18/macOS15起coremltools 引入per_grouped_channel模式由参数group_size指定的一组输出通道共享一个 LUT整个权重矩阵拥有多个LUT。例如形状为(1024, 1024)的权重若group_size16则共有1024 / 16 64个 LUT簇心数量大幅增加聚类近似精度显著提升。源码中的粒度配置与校验粒度在 Core ML 侧由 OpPalettizerConfig 控制关键参数参数默认值说明granularityper_tensor可选per_tensor或per_grouped_channelgroup_size32每组包含的通道数仅per_grouped_channel时生效channel_axisNone分组的通道轴默认None时按 op 类型自动推断nbitsNone每个权重占用的比特数LUT 条目数为2^nbitsweight_threshold2048元素数小于该阈值的权重张量不参与压缩在 PyTorch 侧如 PostTrainingPalettizerConfig 与 ModuleDKMPalettizerConfiggroup_size默认是None并且有严格的互斥校验当granularity per_grouped_channel时必须提供group_sizegroup_size has to be specified along with per_grouped_channel granularity且group_size 0反之per_tensor模式下指定group_size会直接报错。channel_axis默认指向输出通道轴DKM 目前只支持输出通道轴源码校验in_([0])见 palettization_config.py。以(16, 25)形状的权重矩阵为例若group_size8LUT 的形状即为(2, 2^n_bits)——每 8 个通道共享一个含2^n_bits个条目的 LUT。逐通道缩放Per-channel Scale启用enable_per_channel_scaleTrue后权重会先沿输出通道维度按逐通道比例因子per-channel scale归一化再执行聚类。归一化后的权重在相似数值范围内分布更集中可进一步改善聚类质量该模式可与per_grouped_channel组合使用。源码层面的约束值得注意enable_per_channel_scaleTrue与cluster_dim 1向量聚类互斥。在 post_training_palettization.py 与 sensitive_k_means.py 中均有校验违反时会抛出Enabling per_channel_scale is not supported for cluster_dim... larger than 1的断言错误。此外enable_fast_kmeans_mode开启时会将 fp32 权重先转 fp16 再取整以加速聚类该优化同样不与cluster_dim 1兼容。向量聚类Vector Palettizationcluster_dim 详解簇心既可以是标量也可以是向量由参数cluster_dim控制cluster_dim 1默认标量聚类LUT 中每个条目是一个标量cluster_dim 12-D 聚类沿输出通道方向每cluster_dim长度的权重向量共享同一个 2-D 簇心LUT 中每个条目的长度等于cluster_dim。向量聚类能够在更低的位宽下保留相邻元素间的结构信息适合对精度要求更高的场景。Core ML 侧 OpPalettizerConfig 对cluster_dim的描述与 PyTorch 侧ModulePostTrainingPalettizerConfig.cluster_dim的语义完全一致见 post_training_palettization.py。量化 LUT把查找表也压一压LUT 中的值默认以 float 存储。从iOS18/macOS15起LUT 本身也可以按 8-bit 精度存储实现进一步的压缩。Core ML 侧LUT 量化由palettize_weights图 pass 的joint_compression选项及后续量化 pass 支持可将constexpr_lut_to_dense中的 LUT 元素进一步量化见 _quantization_passes.py 中 palettization quantization 的联合压缩说明。PyTorch 侧通过lut_dtype参数控制支持的取值为torch.int8、torch.uint8默认None表示不量化 LUT见 post_training_palettization.py。在 DKM 配置中对应lut_dtype字符串选项i8 / u8 / f16 / bf16 / f32默认f32见 palettization_config.py。可用性Feature AvailabilityCore MLmlprogram模型的 palettized 权重表示自iOS16 / macOS13 / watchOS9 / tvOS16及更新的部署目标deployment target可用per_grouped_channel粒度与 8-bit LUT 则要求iOS18 / macOS15及以上转换时需通过minimum_deployment_targetct.target.iOS18显式声明。三种 Palettization 算法选型coremltools.optimize提供三套算法它们为相同压缩率下的精度表现不同选择时应结合数据、算力与精度诉求算法类型是否需要数据精度/成本权衡K-Means数据无关的 post-training否高比特下精度好低比特下精度损失明显可用per_grouped_channel弥补Sensitive K-MeansSKM基于校准数据的 post-training是少量校准样本对大规模 transformer 架构通常优于数据无关 K-MeansDifferentiable K-MeansDKM训练期training-time是训练数据量要求高压缩-精度权衡最佳可支持极低位宽时间、数据与内存开销最大K-Means无数据后训练聚类最易上手的算法直接用 k-means 聚类权重簇心即 LUT。只需要模型权重无需任何数据。官方精度基准见 opt-palettization-algos.md显示MobileNetv2-1.0 经 6-bit K-Means 后 Top-1 精度 70.89压缩率 2.54×8-bit 为 71.801.97×而 float16 基线为 71.86——可见较高位宽下压缩-精度权衡良好。对应 APIcto.coreml.palettize_weights与cto.torch.palettization.PostTrainingPalettizer。Sensitive K-Means基于 Fisher 信息的加权聚类SKM 源自 SqueezeLLM: Dense-and-Sparse Quantization通过_register_grad_square_hooks在 forward 过程中累计梯度平方。元素越敏感扰动此处即聚类它对模型损失的影响越大加权 k-means 会把簇心向敏感权重值靠近使其表示得更精确从而降低精度劣化。实践中通常128 个校准样本即可SKMPalettizerConfig.calibration_nsamples默认值即为 128见 sensitive_k_means.py且对大型 transformer 架构效果显著优于无数据 K-Means。对应 APIcto.torch.palettization.SKMPalettizer。Differentiable K-Means训练期可微分聚类DKM 是训练期算法每个训练步对权重张量执行可微的软 k-means 聚类分配使流程中每个操作都可微分从而在训练过程中一边更新权重、一边学习 LUT 与索引见 palettizer.py 的类注释。实现方式是在模型内插入模拟聚类的子模块FakePalettize层并用注意力机制计算可微聚类训练结束后通过finalize()将 LUT 与索引折叠回权重。它在所有算法中压缩-精度权衡最佳可支撑 2-bit 等极低位宽官方基准中 ResNet50 经 2-bit DKM 压缩率达 7.63×、精度 75.47 ± 0.05float16 基线 76.14MobileNetv3-small 2-bit 压缩率 5.82×、精度 59.82 ± 0.98。代价是时间与数据开销最大——涉及距离矩阵与注意力矩阵计算内存占用可观。对应 APIcto.torch.palettization.DKMPalettizer。milestone参数可控制从第几步开始启用 palettization默认 0即训练一开始就启用源码中通过_milestones字典跟踪见 palettizer.py。实战压缩一个 Core ML mlprogram 模型基础用法全局 6-bit K-Means对已转换为 Core ML 的模型cto.coreml.palettize_weights会遍历 mlprogram 中所有权重按mode指定的算法构建 LUT将存储权重的constop 替换为constexpr_lut_to_denseop见 _post_training_quantization.py 与 _quantization_passes.py 的图变换说明。下面的例子对所有参数数量超过 512 的 op 应用 6-bit 聚类import coremltools as ct import coremltools.optimize as cto # 加载模型 mlmodel ct.models.MLModel(uncompressed_model_path) # 定义 op 配置6-bit仅压缩元素数超过 512 的权重 op_config cto.coreml.OpPalettizerConfig(nbits6, weight_threshold512) # 将 op 配置全局应用于所有 op config cto.coreml.OptimizationConfig(global_configop_config) # 执行权重聚类压缩 compressed_mlmodel cto.coreml.palettize_weights(mlmodel, config)按 op 类型 / 名称定制OptimizationConfig支持op_type_configs与op_name_configs实现细粒度控制置为None表示跳过该 op。下面的例子对除linear外的全部 op 应用 6-bitlinearop 用 8-bit并跳过名为conv1、conv3的两个卷积import coremltools as ct import coremltools.optimize as cto mlmodel ct.models.MLModel(uncompressed_model_path) global_config cto.coreml.OpPalettizerConfig(nbits6) linear_config cto.coreml.OpPalettizerConfig(nbits8) config cto.coreml.OptimizationConfig( global_configglobal_config, op_type_configs{linear: linear_config}, op_name_configs{conv1: None, conv3: None}, ) compressed_mlmodel cto.coreml.palettize_weights(mlmodel, config)若想按 op 名定制可先用 get_weights_metadata() 工具获取网络中所有权重及其所喂入 op 的详细信息再挑选要定制的 op 名。Core ML 侧 LUT 构建的四种 modeOpPalettizerConfig 通过mode决定 LUT 的构建方式kmeans默认k-means 聚类生成 LUT需提供nbitsuniform线性直方图LUT 为[v_min, v_minscale, ..., v_max]其中scale (v_max - v_min) / (1 nbits - 1)需提供nbitsunique直接用权重中的唯一值构成 LUT权重假设位于离散格点但以 float 存储不能提供nbits系统自动从{1, 2, 4, 6, 8}中取能容纳唯一值数的最小值若唯一值超过 256 个则跳过压缩custom自定义lut_function(weight) - (lut, indices)计算 LUT 与索引不能提供nbits。注意nbits的互斥校验_config.pykmeans/uniform必须提供nbitsunique/custom必须不提供否则抛 ValueError。实战压缩一个 PyTorch 模型Post-Training Palettization数据无关PyTorch 侧流程与 Core ML 侧一致但作用于 Torch 模型先构造配置对象支持代码内字典或从磁盘加载 YAML再调用compress()。下面的例子对全部 op 应用 4-bit并启用per_grouped_channelgroup_size4即每 4 个通道共享一个 LUTfrom model_utilities import get_torch_model from coremltools.optimize.torch.palettization import PostTrainingPalettizer, \ PostTrainingPalettizerConfig # 加载模型 torch_model get_torch_model() palettization_config_dict { global_config: {n_bits: 4, granularity: per_grouped_channel, group_size: 4}, } palettization_config PostTrainingPalettizerConfig.from_dict(palettization_config_dict) palettizer PostTrainingPalettizer(torch_model, palettization_config) palettized_torch_model palettizer.compress()compress()内部把每个受支持模块的参数转成KMeansConfig携带n_bits、axis、lut_dtype、block_size即group_size、cluster_dim、enable_per_channel_scale等随后调用SequentialKMeans或ParallelKMeans完成聚类见 post_training_palettization.py。compress(num_kmeans_workersN)支持多进程并行聚类以加速推荐多 CPU 环境下使用。Sensitive K-MeansYAML 配置 校准数据所有算法的配置都可以写成 YAML 从磁盘加载。以下palettization_config.yaml定义了全局 4-bit、Linear模块 6-bit、跳过conv1/conv3、使用 64 个校准样本global_config: n_bits: 4 module_type_configs: Linear: n_bits: 6 module_name_configs: conv1: null conv3: null calibration_nsamples: 64对应的 Python 脚本加载 YAML 并执行 SKMfrom model_utilities import get_torch_model from data_utils import get_calibration_data from coremltools.optimize.torch.palettization import (SKMPalettizer, SKMPalettizerConfig) import torch.nn.functional as F torch_model get_torch_model() palettization_config SKMPalettizerConfig.from_yaml(palettization_config.yaml) # 定义损失函数前向传播后计算训练损失 loss_fn lambda mod, dat: F.nll_loss(mod(dat[0]), dat[1]) calibration_data get_calibration_data() palettizer SKMPalettizer(torch_model, palettization_config) palettized_torch_model palettizer.compress(dataloadercalibration_data, loss_fnloss_fn)SKM 的compress()比 post-training 版本多两个必填参数dataloader可迭代对象每个元素是模型的一个输入用于计算模型权重的梯度loss_fn可调用对象接收模型与数据、执行前向并计算训练损失。module_name_configs的键支持正则或完全限定名见 sensitive_k_means.py比 Core ML 侧的精确 op 名匹配更灵活。Differentiable K-Means插入训练循环DKM 直接嵌入训练流程一般需要更多数据。标准流程是五个步骤用DKMPalettizerConfig定义聚类参数用DKMPalettizer初始化 palettizer调用prepare(inplaceTrue)将原模型更新为含FakePalettize层的「palettization 友好」模型运行常规训练循环并在每个 step 额外调用palettizer.step()模型收敛后调用finalize()得到可转换到 Core ML 的模型。下面的例子对全部 op 应用 2-bit DKM并指定milestone10——从第 10 个训练步开始启用聚类from model_utilities import get_torch_model from training_utilities import train_step from data_utils import get_dataloader from coremltools.optimize.torch.palettization import (DKMPalettizer, DKMPalettizerConfig) import torch.nn as nn torch_model get_torch_model() dataloader get_dataloader() num_palettization_epochs 2 palettization_config_dict { global_config: {n_bits: 4, milestone: 10}, } palettization_config DKMPalettizerConfig.from_dict(palettization_config_dict) palettizer DKMPalettizer(torch_model, palettization_config) # 向模型中插入 palettization 友好模块 palettizer.prepare(inplaceTrue) torch_model.train() for epoch in range(num_palettization_epochs): for batch_idx, (data, target) in enumerate(dataloader): train_step(data, target, torch_model) palettizer.step() palettized_torch_model palettizer.finalize()ModuleDKMPalettizerConfig还提供大量进阶旋钮绝大多数情况下保持默认即可kmeans_max_iter默认 3、palett_tausoftmax 温度默认 0.0001、enforce_zero是否强制把最接近原点的簇心固定为 0、palett_lambda抑制离群点对簇心计算的影响、palett_unique/palett_shard/palett_batch_mode面向大语言模型的 eDKM 高效选项见 palettization_config.py 的默认值表。若module_type_configs的值是配置列表还会按weight_threshold排序后匹配不同尺寸的权重见 palettization_config.py。将 Palettized PyTorch 模型转换为 Core ML使用ct.optimize.torch.*API 压缩过的 PyTorch 模型coremltools 8 可直接转换无需额外参数但若用到了per_grouped_channel等仅 iOS18/macOS15 可用的新特性必须显式指定部署目标import torch import coremltools as ct palettized_torch_model.eval() traced_model torch.jit.trace(palettized_torch_model, example_input) palettized_coreml_model ct.convert(traced_model, inputs...) # 若聚类时使用了 iOS18 特性 palettized_coreml_model ct.convert(traced_model, inputs.., minimum_deployment_targetct.target.iOS18)如果使用其他非 coremltools 的聚类工具或 coremltools 8 以下版本需参考 opt-conversion.md 中转换带正确 palettized op 的 Core ML 模型的流程。性能与精度取舍Palettization 减小了每个权重值的体积推理时数据搬运量下降对内存瓶颈型memory-bottlenecked模型收益明显该延迟优势仅在权重被「按需即时解压just-in-time decompression」时成立。从 iOS17/macOS14 起这一行为更可能发生在主要运行于 Neural Engine 后端的模型上见 opt-palettization-perf.md。per_grouped_channel模式iOS18/macOS15有一个明确的权衡LUT 数量越多运行时性能越可能下降。经验表明group_size取 8 或 16 通常能在保持精度的同时获得相对未压缩模型的加速。官方基准中 MobileNetv2-1.0 的 2-bit DKM 模型压缩率 5.92×在 iPhone 14 ProiOS17上延迟约 0.47 ms/批与 float16 基线0.48 ms/批基本持平详见 opt-palettization-perf.md 的完整延迟表。测试与验证参考仓库的测试覆盖了本文介绍的三种 APItest_post_training_palettization.py 验证了PostTrainingPalettizer如test_no_config断言无配置时默认 4-bit 聚类后权重唯一值数为 16即2^4个簇心与SKMPalettizertest_palettizer.py 覆盖DKMPalettizertest_palettization_api.py 覆盖配置 API 行为。这些测试是验证配置参数与算法行为的最佳可运行样例也可作为自定义压缩流程的起点。小结Palettization 是 Core ML 生态中与线性量化、剪枝并列的三大模型压缩手段。其核心由「簇心 LUT 索引」组成N{1,2,3,4,6,8}位宽、per_tensor/per_grouped_channel粒度、cluster_dim向量聚类、enable_per_channel_scale逐通道缩放与 8-bit LUT 量化共同构成了完整的配置空间。按数据与算力条件可在无数据 K-Means、基于校准数据的 Sensitive K-Means 与训练期 Differentiable K-Means 之间选择它们分别对应 Core ML 与 PyTorch 两套 API 族。压缩完成后只需在转换时正确声明minimum_deployment_target即可得到可直接部署到 Apple 平台的 mlprogram 模型。赞分享AI 应用机器学习模型优化【免费下载链接】coremltoolsCore ML tools contain supporting tools for Core ML model conversion, editing, and validation.项目地址https://gitcode.com/gh_mirrors/co/coremltools点击查看免费下载相关推荐first-contributions 首次开源贡献实战指南从 Fork 到 Pull Request 的完整 Git 工作流first contributions 首次开源贡献实战指南从 Fork 到 Pull Request 的完整 Git 工作流 本文以 first contrAI 应用机器学习模型优化AutoClip本地视频上传教程讲座、播客、课程如何变高光合集AutoClip本地视频上传教程讲座、播客、课程如何变高光合集 AutoClip 是一款基于 AI 的视频切片与高光生成工具。本教程带你完成 AutoClipAI 应用机器学习模型优化PakePlus终极指南5分钟将网页变桌面应用的完整教程PakePlus终极指南5分钟将网页变桌面应用的完整教程 还在为Electron应用体积庞大、启动缓慢而烦恼吗还在为跨平台打包的环境配置而头疼吗PakeP开发工具桌面应用移动开发跨平台上一篇PaddleOCR 文档视觉语言模型DocVLM模块实战PP-DocBee 系列模型选型、命令行与 Python 推理全指南下一篇tschema性能测试完全指南如何达到每秒500万次验证创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
锦
锦皓数字建站
深耕本土企业品牌数字化升级,专注原创端正雅致商务官网,从视觉设计到稳定运维全程保驾护航。