nrfunc 1.0.0__py3-none-any.whl

nrfunc/__init__.py

@@ -0,0 +1,44 @@
+"""nrfunc：神经区函数算法 —— 成熟模型的「函数化 + 生成物」库。
+
+职责边界（明确，本库只做两件事）：
+
+  上游：你已训练好的成熟模型（任意框架），抽出一层权重矩阵即可。
+  本库：
+    1. 函数化 —— 把权重按「区域函数」重新表达（几何/功能/结构分区 + 0/1 阶函数）
+    2. 生成物 —— 区域函数参数 + 归属索引 + 量化打包 + 分片存储 + 价值对比
+  下游：你自己负责把「生成物」接入推理 / 部署。
+
+本库不做（这些属于上游）：训练、模型定义、蒸馏、微调、量化感知训练（QAT）。
+"""
+from .core import (importance_by_norm, split_regions, kmeans, kmeans_cosine,
+                   assign_rows,
+                   regionify, regionify_hierarchical, reconstruct,
+                   functional_forward, explained_variance,
+                   regionify_transformer_block,
+                   auto_alloc, default_K_grid, default_r_grid)
+from .io import (quantize_weights, dequantize_weights, packed_bytes,
+                 ShardedRegionStore, build_store, shard_placement,
+                 to_bytes, from_bytes)
+from .utils import (bench_single, bench_batch, rss_mb, set_threads, device_of, logf,
+                    size_bytes_raw, size_bytes_regionalized, index_bits,
+                    compare_value, cpu_time, gpu_memory_bytes)
+
+__version__ = '1.0.0'
+
+__all__ = [
+    # 区划分
+    'split_regions', 'importance_by_norm',
+    # 函数化
+    'kmeans', 'kmeans_cosine', 'assign_rows', 'regionify', 'regionify_hierarchical',
+    'reconstruct', 'functional_forward', 'explained_variance',
+    'regionify_transformer_block',
+    'auto_alloc', 'default_K_grid', 'default_r_grid',
+    # 生成物 · 序列化 / 分散存储 / 二进制字节块
+    'quantize_weights', 'dequantize_weights', 'packed_bytes',
+    'ShardedRegionStore', 'build_store', 'shard_placement',
+    'to_bytes', 'from_bytes',
+    # 生成物 · 价值证明 / 性能基准
+    'bench_single', 'bench_batch', 'rss_mb', 'set_threads', 'device_of', 'logf',
+    'size_bytes_raw', 'size_bytes_regionalized', 'index_bits',
+    'compare_value', 'cpu_time', 'gpu_memory_bytes',
+]

nrfunc/_rand.py

@@ -0,0 +1,35 @@
+"""确定性随机（内置 PRNG，保证可复现）。
+
+库内所有需要随机的环节（聚类初始质心、种子派生）统一走 mulberry32 /
+fisher_yates，保证同一 seed 跨次运行、跨平台结果一致。
+"""
+import math
+
+
+def mulberry32(seed):
+    """mulberry32 确定性 PRNG，返回无参函数 ``rand() -> [0, 1)``。
+
+    保证同一 seed 跨次运行、跨平台结果一致。
+    """
+    a = seed & 0xFFFFFFFF
+
+    def rand():
+        nonlocal a
+        a = (a + 0x6D2B79F5) & 0xFFFFFFFF
+        x = (a ^ (a >> 15)) & 0xFFFFFFFF
+        t = (x * (1 | a)) & 0xFFFFFFFF
+        y = (t ^ (t >> 7)) & 0xFFFFFFFF
+        t = ((t + (y * (61 | t)) & 0xFFFFFFFF) ^ t) & 0xFFFFFFFF
+        t = (t ^ (t >> 14)) & 0xFFFFFFFF
+        return t / 4294967296.0
+
+    return rand
+
+
+def fisher_yates(n, rand):
+    """Fisher-Yates 洗牌，返回打乱后的下标列表 [0, n)。"""
+    idx = list(range(n))
+    for i in range(n - 1, 0, -1):
+        j = int(math.floor(rand() * (i + 1)))
+        idx[i], idx[j] = idx[j], idx[i]
+    return idx

nrfunc/core/__init__.py

@@ -0,0 +1,15 @@
+"""core：区划分 + 函数化（神经区函数算法的核心算子）。"""
+from .region import importance_by_norm, split_regions
+from .functionalize import (kmeans, kmeans_cosine, assign_rows, regionify,
+                            regionify_hierarchical, reconstruct,
+                            functional_forward, explained_variance,
+                            regionify_transformer_block)
+from .auto_alloc import auto_alloc, default_K_grid, default_r_grid
+
+__all__ = [
+    'importance_by_norm', 'split_regions',
+    'kmeans', 'kmeans_cosine', 'assign_rows', 'regionify', 'regionify_hierarchical',
+    'reconstruct', 'functional_forward', 'explained_variance',
+    'regionify_transformer_block',
+    'auto_alloc', 'default_K_grid', 'default_r_grid',
+]

nrfunc/core/auto_alloc.py

@@ -0,0 +1,266 @@
+"""自适应函数化分配器（auto_alloc）：给每一层自动挑「函数化配置」，或判「跳过」。
+
+这是 nrfunc「函数化」能力的自适应前端：不训练、不蒸馏、不碰 QAT，只对已训练好的
+权重，在「内存净省」+「掉点可容忍」两个约束下，自动决定这一层：
+  * 用不用函数化（不划算就跳过）
+  * 用 0 阶（质心）还是 1 阶（低秩）
+  * K（分区数）取多少、r（主成分数）取多少
+
+为什么需要它（P3 v0.8 立项动机）：
+  1 阶低秩「净省内存」是有条件的（近似充要条件 K×(1+r) < n）。给小层 / 深层
+  无脑上同一套 K、r 会内存反涨（实测反例：小 MLP fc1 多 158%、小 CNN conv1
+  多 73%、大 CNN 最深卷积层多 25%）。auto_alloc 逐层搜索，只保留「净省」的
+  配置，从机制上杜绝反涨。
+"""
+from __future__ import annotations
+
+import math
+
+import numpy as np
+
+from .functionalize import assign_rows, regionify
+from ..utils import size_bytes_raw, size_bytes_regionalized, index_bits
+
+__all__ = ['auto_alloc', 'default_K_grid', 'default_r_grid']
+
+
+def _size_after_est(n, D, K, r, order, bits):
+    """纯解析式预估函数化后字节数（无需跑 regionify）。
+
+    由 size_bytes_regionalized 的公式直接推出：regionify 结果中 components
+    必为 (K, r)、coeffs 必为 (n, r)，故函数化后字节数只由 (n, D, K, r, order)
+    决定，可在跑昂贵聚类/SVD 之前精确预估，用作净省先验剪枝。
+    """
+    ibits = index_bits(n, K)
+    if order == 0:
+        fn_params = K * D
+    else:
+        fn_params = K * D + K * r * D + n * r
+    return math.ceil(fn_params * bits / 8) + math.ceil(n * ibits / 8)
+
+
+def default_K_grid(n):
+    """默认 K（分区数）候选网格：2 的幂（≤ n），并补一个 sqrt(n) 附近参考值。"""
+    if n < 2:
+        return [1]
+    grid = []
+    k = 2
+    while k <= n:
+        grid.append(k)
+        k *= 2
+    s = int(round(np.sqrt(n)))
+    if s >= 2 and s not in grid:
+        grid.append(s)
+    return sorted(set(grid))
+
+
+def default_r_grid():
+    """默认 r（1 阶每区主成分数）候选网格。"""
+    return [1, 2, 4, 8]
+
+
+def auto_alloc(rows, signal='G', activations=None, groups=None,
+               budget=None, tol_ev=None, eval_fn=None,
+               fidelity_margin=0.3,
+               K_grid=None, r_grid=None, r_max=None,
+               iters=60, seed=0, float_bits=32):
+    """给一层权重自动挑最优函数化配置，或判「跳过」（保持原样）。
+
+    参数：
+      rows：(n, D) 被函数化的权重行向量（n=单元数，D=每单元参数量）。
+      signal：'G' 几何 / 'F' 功能（需 activations）/ 'S' 结构（需 groups）。
+      activations：signal='F' 时 (n, T) 激活响应矩阵。
+      groups：signal='S' 时 (n,) 每个单元的结构归属（层号/头号/通道号）。
+      budget：可选，函数化后字节上限。给定时该层目标从「净省」收紧为
+        「after ≤ budget」；若原权重 before ≤ budget 则该层直接判跳过。
+      tol_ev：可选，重构保真度（explained_variance）下限。注意：explained 是
+        「重构误差」代理，跨信号不可比（S 结构分区 explained 天然偏低但分类
+        掉点低），故默认不启用；掉点约束更可靠的手段是 eval_fn。
+      eval_fn：可选，真实精度回调 eval_fn(result) -> bool，True=接受该候选。
+        这是「掉点 < 阈值」的可靠实现（库不内置数据/前向，交给使用方）。
+      fidelity_margin：默认 0.3。相对保真保护——在「净省」候选中，排除重构
+        保真度（explained）比同层最优候选低 0.3 以上的，防止为多省一点内存而
+        把保真砸穿。S 结构分区 explained 天然偏低、相对差距小，故不误伤；
+        传 None 关闭（退化为纯内存口径「省最多」）。
+      K_grid / r_grid：可选候选网格。默认 K=2 的幂≤n（补 sqrt(n)），r=[1,2,4,8]。
+      r_max：可选，1 阶主成分数上限（对 r_grid 截断）。
+      iters / seed：聚类迭代与随机种子（固定 seed 保证可复现）。
+      float_bits：内存口径位宽，默认 32（f32 常驻，与「内存·函数部署」口径一致）。
+
+    返回 dict：
+      plan：(order, K, r) 三元组；None 表示「跳过」（该层保持原样，最诚实）。
+      decision：'order1' / 'order0' / 'skip'
+      result：regionify 返回 dict（跳过时为 None）
+      order / K / r：选中的阶数与参数（跳过时均为 None）
+      bytes_before / bytes_after：函数化前后内存字节（float_bits 口径）
+      saving：节省率（正=省，负=多；跳过=0.0）
+      explained：重构保真度（跳过时为 None）
+      n_candidates：通过全部约束的候选数
+      n_tried：实际跑过 regionify 的候选总数
+      skipped_reason：跳过原因（仅 decision='skip' 时非空）
+    """
+    rows = np.asarray(rows, dtype=np.float64)
+    n, D = rows.shape
+    bytes_before = size_bytes_raw(n, D, float_bits)
+
+    # —— 预算语义：原权重已达标则无需函数化（省事且保精度）——
+    if budget is not None and bytes_before <= budget:
+        return _skip(bytes_before, 0, '原权重已满足内存预算（before ≤ budget）')
+
+    # —— 候选网格 ——
+    if signal == 'S':
+        groups_arr = np.asarray(groups)
+        if groups_arr.size == 0:
+            raise ValueError("signal='S'（结构分区）需提供非空 groups")
+        K_vals = [int(len(np.unique(groups_arr)))]  # 结构边界固定 K，不搜
+    else:
+        K_vals = [int(k) for k in (K_grid if K_grid is not None else default_K_grid(n))]
+
+    r_vals = [int(r) for r in (r_grid if r_grid is not None else default_r_grid())]
+    if r_max is not None:
+        r_vals = [r for r in r_vals if r <= r_max]
+
+    # —— 逐候选构造模型。同一 K 的分区只聚类一次：regionify 的聚类阶段对
+    # 同 K、同 seed 结果逐位相同，重复跑是纯浪费（原来每 K 跑 1+len(r_vals) 次）——
+    candidates = []
+    n_tried = 0
+    assign_cache = {}
+    for order in (1, 0):
+        for K in K_vals:
+            if K > n:
+                continue
+            if order == 0:
+                if K >= n:  # K=n 退化为逐单元无压缩，必然不省
+                    continue
+                # 精确净省预判：预估字节 ≥ 原字节则必然不省，跳过（省去聚类）
+                if _size_after_est(n, D, K, None, 0, float_bits) >= bytes_before:
+                    continue
+                assign = _get_assign(assign_cache, rows, K, signal,
+                                     iters, seed, activations, groups)
+                res = regionify(rows, signal=signal, K=K, order=0,
+                                iters=iters, seed=seed,
+                                activations=activations, groups=groups,
+                                assign=assign)
+                n_tried += 1
+                candidates.append(_candidate(res, order, K, None,
+                                             bytes_before, float_bits))
+            else:  # order == 1
+                for r in r_vals:
+                    if r < 1:
+                        continue
+                    # 先验剪枝：1 阶净省的必要条件是 K×(1+r) < n（否则必然反涨）
+                    if K * (1 + r) >= n:
+                        continue
+                    # 精确净省预判：预估字节 ≥ 原字节则必然不省，跳过（省去 SVD）
+                    if _size_after_est(n, D, K, r, 1, float_bits) >= bytes_before:
+                        continue
+                    assign = _get_assign(assign_cache, rows, K, signal,
+                                         iters, seed, activations, groups)
+                    res = regionify(rows, signal=signal, K=K, order=1, r=r,
+                                    iters=iters, seed=seed,
+                                    activations=activations, groups=groups,
+                                    assign=assign)
+                    n_tried += 1
+                    candidates.append(_candidate(res, order, K, r,
+                                                 bytes_before, float_bits))
+
+    # —— 硬约束：内存净省（绝不反涨）——
+    kept = [c for c in candidates if c['bytes_after'] < bytes_before]
+    if not kept:
+        return _skip(bytes_before, n_tried,
+                     '候选网格内不存在任何净省配置（该层冗余度不足，函数化必反涨）')
+
+    # —— 可选：内存预算 ——
+    if budget is not None:
+        kept = [c for c in kept if c['bytes_after'] <= budget]
+        if not kept:
+            return _skip(bytes_before, n_tried, '存在净省候选，但均超出内存预算 budget')
+
+    # —— 掉点代理：默认相对保真保护 + 可选绝对 tol_ev ——
+    if fidelity_margin is not None:
+        ev_max = max(c['explained'] for c in kept)
+        kept = [c for c in kept if c['explained'] >= ev_max - fidelity_margin]
+        if not kept:
+            return _skip(bytes_before, n_tried,
+                         f'存在净省候选，但均未通过相对保真保护（fidelity_margin={fidelity_margin}）')
+    if tol_ev is not None:
+        kept = [c for c in kept if c['explained'] >= tol_ev]
+        if not kept:
+            return _skip(bytes_before, n_tried,
+                         f'存在净省候选，但重构保真度均低于 tol_ev={tol_ev}')
+
+    # —— 真实精度（可选，最可靠）——
+    # 早停优化：最终选「省最多、并列取保真最高」，故按 (saving, explained) 降序
+    # 逐个评估，第一个通过 eval_fn 的即全局最优，其余候选无需再跑昂贵精度评估。
+    if eval_fn is not None:
+        ordered = sorted(kept, key=lambda c: (-c['saving'], -c['explained']))
+        passed = None
+        for c in ordered:
+            if bool(eval_fn(c['result'])):
+                passed = c
+                break
+        if passed is None:
+            return _skip(bytes_before, n_tried, '存在净省候选，但均未通过 eval_fn 真实精度约束')
+        best = passed
+    else:
+        # 选「省最多」，并列取「保真最高」（explained 越大 = 掉点越小）
+        best = max(kept, key=lambda c: (c['saving'], c['explained']))
+
+    return {
+        'plan': (best['order'], best['K'], best['r']),
+        'decision': 'order1' if best['order'] == 1 else 'order0',
+        'result': best['result'],
+        'order': best['order'],
+        'K': best['K'],
+        'r': best['r'],
+        'bytes_before': bytes_before,
+        'bytes_after': best['bytes_after'],
+        'saving': best['saving'],
+        'explained': best['explained'],
+        'n_candidates': len(kept),
+        'n_tried': n_tried,
+        'skipped_reason': None,
+    }
+
+
+def _get_assign(cache, rows, K, signal, iters, seed, activations, groups):
+    """按 K 取缓存的聚类归属；未缓存才跑一次聚类（同 K 全程只跑一次）。"""
+    assign = cache.get(K)
+    if assign is None:
+        assign, _ = assign_rows(rows, signal=signal, K=K, iters=iters, seed=seed,
+                                activations=activations, groups=groups)
+        cache[K] = assign
+    return assign
+
+
+def _candidate(result, order, K, r, bytes_before, float_bits):
+    """把一次 regionify 结果封装成候选，算好字节与节省率。"""
+    after = size_bytes_regionalized(result, bits=float_bits)
+    saving = round((bytes_before - after) / bytes_before, 4)
+    return {
+        'result': result,
+        'order': order,
+        'K': K,
+        'r': r,
+        'bytes_after': after,
+        'saving': saving,
+        'explained': result.get('explained'),
+    }
+
+
+def _skip(bytes_before, n_tried, reason):
+    return {
+        'plan': None,
+        'decision': 'skip',
+        'result': None,
+        'order': None,
+        'K': None,
+        'r': None,
+        'bytes_before': bytes_before,
+        'bytes_after': None,
+        'saving': 0.0,
+        'explained': None,
+        'n_candidates': 0,
+        'n_tried': n_tried,
+        'skipped_reason': reason,
+    }