"""examples 共享工具：路径引导 + MNIST 图像分类数据加载 + 训练/评估助手。
用法（各示例文件顶部）：
from _common import (load_mnist, train_model, accuracy, device,
downsample_images)
说明：
- 路径引导：未 `pip install -e .` 时，也能从本仓库直接 `import nrfunc`。
- MNIST 加载：优先复用本仓库已有缓存（prototypes/neural-region-p1/data/*.bin，
与官方 MNIST 同格式），没有则下载官方 MNIST 到 examples/_data/（幂等）。
- 数据格式（.bin，自定轻量缓存，与官方 idx 等价）：
uint32 n（大端） + float32[n×784]（小端，像素∈[0,1]） + uint8[n]（标签）
"""
import gzip
import os
import sys
import urllib.request
_CPU_CAP = int(os.environ.get('NRFUNC_CPU_THREADS', '0') or 0)
if _CPU_CAP <= 0:
_CPU_CAP = max(2, min(os.cpu_count() or 4, 8))
for _k in ('OMP_NUM_THREADS', 'MKL_NUM_THREADS', 'OPENBLAS_NUM_THREADS',
'NUMEXPR_NUM_THREADS'):
os.environ.setdefault(_k, str(_CPU_CAP))
import numpy as np
_HERE = os.path.dirname(os.path.abspath(__file__))
_SRC = os.path.abspath(os.path.join(_HERE, '..', 'src'))
if _SRC not in sys.path:
sys.path.insert(0, _SRC)
_DATA_DIR = os.path.join(_HERE, '_data')
_PROTO_CACHE = os.path.abspath(
os.path.join(_HERE, '..', '..', 'prototypes', 'neural-region-p1', 'data'))
_BASE = 'https://ossci-datasets.s3.amazonaws.com/mnist/'
_FILES = [
'train-images-idx3-ubyte.gz', 'train-labels-idx1-ubyte.gz',
't10k-images-idx3-ubyte.gz', 't10k-labels-idx1-ubyte.gz',
]
_DIM = 784
def _read_bin(path):
raw = np.fromfile(path, dtype=np.uint8)
n = int.from_bytes(raw[:4].tobytes(), 'big')
px = raw[4:4 + n * _DIM * 4].view(np.float32).reshape(n, _DIM)
lab = raw[4 + n * _DIM * 4:4 + n * _DIM * 4 + n].astype(np.int64)
return px, lab
def _download():
os.makedirs(_DATA_DIR, exist_ok=True)
print('[数据] 未找到本地缓存，下载官方 MNIST ...')
gzs = {}
for name in _FILES:
url = _BASE + name
print(f' 下载 {url}')
with urllib.request.urlopen(url, timeout=120) as res:
gzs[name] = res.read()
def _parse_images(buf):
raw = gzip.decompress(buf)
n = int.from_bytes(raw[4:8], 'big')
dim = int.from_bytes(raw[8:12], 'big') * int.from_bytes(raw[12:16], 'big')
px = np.frombuffer(raw, dtype=np.uint8, offset=16,
count=n * dim).astype(np.float32) / 255.0
return n, dim, px
def _parse_labels(buf):
raw = gzip.decompress(buf)
n = int.from_bytes(raw[4:8], 'big')
return np.frombuffer(raw, dtype=np.uint8, offset=8,
count=n).astype(np.int64)
for kind, img, lab in (
('train', 'train-images-idx3-ubyte.gz', 'train-labels-idx1-ubyte.gz'),
('test', 't10k-images-idx3-ubyte.gz', 't10k-labels-idx1-ubyte.gz')):
n, dim, px = _parse_images(gzs[img])
y = _parse_labels(gzs[lab])
buf = np.empty(4 + n * dim * 4 + n, dtype=np.uint8)
buf[:4] = np.array([n], dtype='>u4').view(np.uint8)
buf[4:4 + n * dim * 4] = px.astype('<f4').view(np.uint8)
buf[4 + n * dim * 4:] = y.astype(np.uint8)
buf.tofile(os.path.join(_DATA_DIR, kind + '.bin'))
print('[数据] 下载完成。')
def load_mnist():
"""加载 MNIST，返回 (train_x, train_y, test_x, test_y)。x 为 float32∈[0,1]，y 为 int64。"""
own_train = os.path.join(_DATA_DIR, 'train.bin')
own_test = os.path.join(_DATA_DIR, 'test.bin')
proto_train = os.path.join(_PROTO_CACHE, 'train.bin')
proto_test = os.path.join(_PROTO_CACHE, 'test.bin')
train_bin = own_train if os.path.exists(own_train) else (
proto_train if os.path.exists(proto_train) else None)
test_bin = own_test if os.path.exists(own_test) else (
proto_test if os.path.exists(proto_test) else None)
if train_bin is None or test_bin is None:
_download()
train_bin, test_bin = own_train, own_test
return (_read_bin(train_bin)[0], _read_bin(train_bin)[1],
_read_bin(test_bin)[0], _read_bin(test_bin)[1])
def device():
import torch
try:
torch.set_num_threads(_CPU_CAP)
except Exception:
pass
return 'cuda' if torch.cuda.is_available() else 'cpu'
def downsample_images(x, size=7):
"""把 28×28 图像按块平均缩到 size×size，返回 [N, size*size] float32。
size 须整除 28（可选 4 / 7 / 14）。
"""
n = x.shape[0]
step = 28 // size
assert step * size == 28, 'size 须整除 28'
x = x.reshape(n, 28, 28)
x = x.reshape(n, size, step, size, step).mean(axis=(2, 4))
return x.reshape(n, size * size)
def train_model(model, x, y, epochs=15, batch=256, lr=1e-3, seed=0,
dev=None, verbose=True):
"""通用小网络训练（交叉熵），返回训练后的 model（已 eval）。
model 须返回 logits（最后一维 = 类别数）。
"""
import torch
if dev is None:
dev = device()
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
model.apply(lambda m: m.reset_parameters()
if hasattr(m, 'reset_parameters') else None)
model = model.to(dev)
xt = torch.as_tensor(x, dtype=torch.float32, device=dev)
yt = torch.as_tensor(y, dtype=torch.long, device=dev)
n = len(xt)
opt = torch.optim.Adam(model.parameters(), lr=lr)
lossf = torch.nn.CrossEntropyLoss()
for e in range(1, epochs + 1):
perm = torch.randperm(n, device=dev)
model.train()
total, cnt = 0.0, 0
for b in range(0, n, batch):
ids = perm[b:b + batch]
opt.zero_grad()
loss = lossf(model(xt[ids]), yt[ids])
loss.backward()
opt.step()
total += loss.item() * len(ids)
cnt += len(ids)
if verbose and (e % 5 == 0 or e == epochs):
acc = accuracy(model, xt[:8192], yt[:8192], dev)
print(f' epoch {e:2d}/{epochs} loss={total / cnt:.4f} acc={acc * 100:.2f}%')
model.eval()
return model
def accuracy(model, x, y, dev=None, batch=1024):
import torch
if dev is None:
dev = device()
model.eval()
with torch.no_grad():
xt = torch.as_tensor(x, dtype=torch.float32, device=dev)
yt = torch.as_tensor(y, dtype=torch.long, device=dev)
n = len(xt)
correct = 0
for b in range(0, n, batch):
pred = model(xt[b:b + batch]).argmax(dim=1)
correct += (pred == yt[b:b + batch]).sum().item()
return correct / n
def show_artifacts(res, bits=8, n_shards=4):
"""展示「生成物」环节（库的核心产出）：量化打包 + 分片存储 + 并行读取。
对应完整链路「训练 → 模型 → 函数化 → 生成物 → 部署使用」里的「生成物」。
函数化产物 res 在这里被打包成可部署字节、分散到多节点并可并行读回。
"""
import nrfunc
shared = res['centroids'] if res['order'] == 0 else res['means']
packed, scale, n_out, n_in = nrfunc.quantize_weights(shared, bits=bits)
store = nrfunc.build_store(res, n_shards=n_shards)
ids = list(range(min(res['K'], n_shards)))
params = store.region_read_many(ids)
kind = 'centroids' if res['order'] == 0 else 'means'
print(f' · 量化打包：区域函数 {kind} → {len(packed)} 字节（{bits}bit）')
print(f' · 分片存储：{n_shards} 节点，分布={store.placement()}，'
f'负载={store.load_balance().tolist()}')
print(f' · 并行读取：读回 {len(params)} 个区域函数')
def _compact_tensors(res):
"""函数化后的紧凑参数张量列表（用于 GPU 显存实测的「后」口径）。"""
if res['order'] == 0:
return [res['centroids']]
tensors = [res['means'], res['components']]
if res.get('coeffs') is not None:
tensors.append(res['coeffs'])
return tensors
def make_drop_guard(model, layer, orig_shape, x, y, dev, acc_before, tol_pt=3.0):
"""构造 auto_alloc 的 eval_fn：真实精度回调，掉点 < tol_pt 才接受该候选。
auto_alloc 会对每个「净省 + 保真」候选调用 eval_fn(result) -> bool，这里把
候选 reconstruct 回 layer 权重、测整网准确率、判断相对 acc_before 的掉点是否
低于 tol_pt，并在每次测完恢复 layer 原权重（避免污染后续候选的评估）。
这是「掉点 < 3pt 红线」的可靠实现（库不内置数据/前向，交给使用方）；相比库
默认的 explained 重构保真代理，它能拦住「首层 0 阶质心灾难掉点」这类代理
看不见的坑（如小 MLP 首层 0 阶会把 88% 砸到 27%）。
"""
import torch
import nrfunc
backup = layer.weight.detach().clone()
def _guard(result):
if result is None:
return False
recon = nrfunc.reconstruct(result)
with torch.no_grad():
layer.weight.copy_(
torch.as_tensor(recon.reshape(orig_shape),
dtype=layer.weight.dtype, device=layer.weight.device))
acc = accuracy(model, x, y, dev)
with torch.no_grad():
layer.weight.copy_(backup)
return (acc_before - acc) * 100 < tol_pt
return _guard
def functionalize_layers(model, layers, test_x, test_y, dev, acc_before,
tol_pt=3.0, seed=0):
"""逐层函数化（参差式）：对多个层各自独立 auto_alloc，并叠加测整网掉点。
这是「一层层分别函数化」的完整流程，分两步：
1. 逐层独立决策：每层用 make_drop_guard 的真实精度掉点作 eval_fn，
各自选出最优 (order,K,r)（不划算的层判 skip、保持原样）。
2. 叠加部署：把已决策的各层「同时」替换成重建权重，测整网准确率，
得到误差累积后的真实掉点（略大于单层掉点之和）。
参数：
model / layers：模型及其要函数化的层列表 [(name, layer), ...]
test_x, test_y, dev, acc_before：评估数据/设备/原始准确率
tol_pt：掉点红线（默认 3pt）
seed：auto_alloc 随机种子（可复现）
返回 dict：
layers：[{name, shape, decision, order, K, r, saving, explained,
single_drop, res}] 每层决策与单层掉点
acc_before / acc_after：叠加部署前后整网准确率
stacked_drop：整网叠加掉点（pt）
total_before / total_after / total_saving：三层（或 N 层）合计内存字节与节省率
诚实口径：单层掉点（single_drop）是「只换该层」的独立性值；叠加掉点
（stacked_drop）是「所有层同时换」的真实上线值，约 ≥ 单层掉点之和（误差累积）。
"""
import torch
import nrfunc
layer_list = [(name, layer) for name, layer in layers]
originals = {name: layer.weight.detach().clone() for name, layer in layer_list}
per_layer = []
for name, layer in layer_list:
W = layer.weight.detach().cpu().numpy().astype(np.float64)
shape = W.shape
guard = make_drop_guard(model, layer, shape, test_x, test_y,
dev, acc_before, tol_pt=tol_pt)
alloc = nrfunc.auto_alloc(W, signal='G', eval_fn=guard, seed=seed)
if alloc['decision'] == 'skip':
single_drop = None
res = None
else:
res = alloc['result']
recon = nrfunc.reconstruct(res)
with torch.no_grad():
layer.weight.copy_(torch.as_tensor(
recon.reshape(shape), dtype=layer.weight.dtype,
device=layer.weight.device))
acc_layer = accuracy(model, test_x, test_y, dev)
single_drop = (acc_before - acc_layer) * 100
with torch.no_grad():
layer.weight.copy_(originals[name])
per_layer.append({
'name': name, 'shape': shape,
'decision': alloc['decision'],
'order': alloc['order'], 'K': alloc['K'], 'r': alloc['r'],
'saving': alloc['saving'], 'explained': alloc['explained'],
'single_drop': single_drop, 'res': res,
'bytes_before': alloc['bytes_before'],
'bytes_after': alloc['bytes_after'],
})
with torch.no_grad():
for item, (name, layer) in zip(per_layer, layer_list):
if item['res'] is None:
layer.weight.copy_(originals[name])
else:
recon = nrfunc.reconstruct(item['res'])
layer.weight.copy_(torch.as_tensor(
recon.reshape(item['shape']), dtype=layer.weight.dtype,
device=layer.weight.device))
acc_after = accuracy(model, test_x, test_y, dev)
stacked_drop = (acc_before - acc_after) * 100
with torch.no_grad():
for name, layer in layer_list:
layer.weight.copy_(originals[name])
total_before = sum(i['bytes_before'] for i in per_layer)
total_after = sum(i['bytes_after'] if i['bytes_after'] is not None
else i['bytes_before'] for i in per_layer)
return {
'layers': per_layer,
'acc_before': acc_before, 'acc_after': acc_after,
'stacked_drop': stacked_drop,
'total_before': total_before, 'total_after': total_after,
'total_saving': round((total_before - total_after) / total_before, 4),
'n_layers': len(per_layer),
'n_skip': sum(1 for i in per_layer if i['decision'] == 'skip'),
}
def show_value(W, res, x_batch, bits=8, dev=None, float_bits=32):
"""价值证明 · 完整四维：体积 / 内存 / CPU / GPU，函数化前 vs 后。
float_bits：运行时常驻浮点精度位宽（32 或 64），由使用方按权重精度要求自由选择。
默认 32（通用小模型）；对权重精度敏感的大模型（如大 CNN）可传 64。
口径（诚实，不夸大）：
- 体积：权重存储字节（float_bits → bits 量化），对应「存储 / 传输」占用。
- 内存：运行时常驻字节，分两种部署方式——
还原部署（reconstruct 回 float 稠密，精度换体积、内存不省）；
函数部署（functional_forward 不还原，只常驻紧凑参数，体积+内存都省）。
- CPU ：同机同批实测。函数化前=稠密前向 x@W^T；函数化后=函数前向
functional_forward（不还原）。含单样本延迟(ns) + 批量吞吐(samples/s)。
- GPU ：同机同批实测（需 CUDA）。与 CPU 侧对称，含：
显存（前=原始权重 vs 后=紧凑参数搬上显存）；
单样本延迟 + 批量吞吐（前=稠密前向 vs 后=函数前向，torch 真实 CUDA 前向）。
注：GPU 函数前向需 Rust/CUDA kernel 沉淀；此处的 torch 实现在是
「同一 CUDA 环境下稠密 vs 函数」的可比口径。
"""
import nrfunc
import numpy as np
import torch
_dtype = np.float64 if float_bits >= 64 else np.float32
W = np.asarray(W, dtype=_dtype)
x_batch = np.asarray(x_batch, dtype=_dtype)
x1 = x_batch[:1]
n, D = W.shape
def _dense(x):
return x @ W.T
def _func(x):
return nrfunc.functional_forward(res, x)
cpu_before_ns = nrfunc.cpu_time(_dense, x1, iters=500, warmup=20)
cpu_after_ns = nrfunc.cpu_time(_func, x1, iters=500, warmup=20)
B = int(len(x_batch))
sps_cpu_before = B / (nrfunc.cpu_time(_dense, x_batch, iters=50, warmup=10) * 1e-9)
sps_cpu_after = B / (nrfunc.cpu_time(_func, x_batch, iters=50, warmup=10) * 1e-9)
gpu_before_mb = gpu_after_mb = None
gpu_ns_before = gpu_ns_after = None
sps_gpu_before = sps_gpu_after = None
if dev is None:
dev = 'cuda' if torch.cuda.is_available() else 'cpu'
if dev == 'cuda' and torch.cuda.is_available():
try:
def _alloc_mb(tensors):
arrs = [np.ascontiguousarray(np.asarray(t, dtype=_dtype))
for t in tensors]
torch.zeros(1, device='cuda')
torch.cuda.synchronize()
base = torch.cuda.memory_allocated()
ts = [torch.as_tensor(a, device='cuda') for a in arrs]
torch.cuda.synchronize()
alloc = torch.cuda.memory_allocated() - base
del ts
return alloc / 1048576.0
gpu_before_mb = _alloc_mb([W])
gpu_after_mb = _alloc_mb(_compact_tensors(res))
Wc = torch.as_tensor(np.ascontiguousarray(W), device='cuda')
x1c = torch.as_tensor(np.ascontiguousarray(x1), device='cuda')
xbc = torch.as_tensor(np.ascontiguousarray(x_batch), device='cuda')
def _dense_gpu(x):
return x @ Wc.T
def _func_gpu(x):
return _gpu_forward(res, x)
def _gpu_time(fn, arg, iters=100, warmup=20):
for _ in range(warmup):
fn(arg)
torch.cuda.synchronize()
s = torch.cuda.Event(enable_timing=True)
e = torch.cuda.Event(enable_timing=True)
s.record()
for _ in range(iters):
fn(arg)
e.record()
torch.cuda.synchronize()
return s.elapsed_time(e) / iters * 1e6
gpu_ns_before = _gpu_time(_dense_gpu, x1c)
gpu_ns_after = _gpu_time(_func_gpu, x1c)
sps_gpu_before = B / (_gpu_time(_dense_gpu, xbc, iters=20, warmup=5) * 1e-9)
sps_gpu_after = B / (_gpu_time(_func_gpu, xbc, iters=20, warmup=5) * 1e-9)
del Wc, x1c, xbc
except Exception:
pass
vol_before = nrfunc.size_bytes_raw(n, D, float_bits)
vol_after = nrfunc.size_bytes_regionalized(res, bits=bits)
mem_before = n * D * (float_bits // 8)
mem_func_after = nrfunc.size_bytes_regionalized(res, bits=float_bits)
def _save(b, a, higher_better=False):
if b == 0:
return 0.0
return (a - b) / b if higher_better else (b - a) / b
ftag = f'f{float_bits}'
dims = [
(f'体积(存储,{float_bits}→{bits}bit)', vol_before, vol_after, False),
(f'内存·还原部署({ftag})', mem_before, mem_before, False),
(f'内存·函数部署({ftag})', mem_before, mem_func_after, False),
('CPU单样本延迟(ns)', cpu_before_ns, cpu_after_ns, False),
('CPU批量吞吐(samples/s)', sps_cpu_before, sps_cpu_after, True),
]
if gpu_before_mb is not None:
dims.append(('GPU显存(MB)', gpu_before_mb, gpu_after_mb, False))
if gpu_ns_before is not None:
dims.append(('GPU单样本延迟(ns)', gpu_ns_before, gpu_ns_after, False))
if sps_gpu_before is not None:
dims.append(('GPU批量吞吐(samples/s)', sps_gpu_before, sps_gpu_after, True))
lines = [f'{"维度":<24}{"函数化前":>16}{"函数化后":>16}{"节省率":>12}', '-' * 68]
for name, b, a, hb in dims:
lines.append(f'{name:<24}{b:>16,.1f}{a:>16,.1f}{_save(b, a, hb):>11.1%}')
table = '\n'.join(lines)
print(table)
print('\n口径说明（诚实）：')
print(f' · 体积省得最多（{bits}bit 量化）；内存/显存是 {ftag} 常驻口径，省得少些。')
print(' · 内存分「还原部署」（reconstruct 回稠密，内存不省）与「函数部署」')
print(' （functional_forward 不还原，只常驻紧凑参数）两种，选后者才省内存。')
print(' · CPU 单样本/批量：numpy 参考实现的 functional_forward 在 1 阶小模型下')
print(' 受 Python 循环+gather 限制，可能慢于 BLAS 稠密；提速只在 K远小于n 的大模型')
print(' + 编译语言/GPU 里兑现（见 functional_forward_demo.py）。')
gpu_lines = []
if gpu_before_mb is not None:
gpu_lines.append(f'显存={ftag} 常驻（同「内存·函数部署」）')
if gpu_ns_before is not None:
gpu_lines.append('单样本延迟 + 批量吞吐=torch 真实 CUDA 前向（稠密 vs 函数）')
if gpu_lines:
print(' · GPU：' + '；'.join(gpu_lines) + '。')
print(' · 说明：GPU 函数前向的提速潜力仍需 Rust/CUDA kernel 沉淀，此处 torch 实测')
print(' 是「同一 CUDA 环境」下的可比口径，不代表已兑现生产级 GPU 提速。')
return {'dims': dims, 'table': table, 'explained': res.get('explained')}
def _gpu_forward(result, x):
"""torch 版函数前向：在 CUDA 张量上重演 functional_forward 的语义。
与 numpy 的 functional_forward 数值一致（float 误差内），但输入/权重均在
显存上，用于 GPU 侧耗时实测。仅覆盖 0 阶 / 1 阶；1 阶低秩修正走「逐区 GEMM」
（与 numpy 大批量分支同语义），避免构造 (n, K*r) 稠密展开撑爆显存。
"""
import torch
dtype = x.dtype
assign = torch.as_tensor(result['assign'], device='cuda')
if result['order'] == 0:
C = torch.as_tensor(np.ascontiguousarray(result['centroids']),
device='cuda', dtype=dtype)
z = x @ C.T
return z.gather(-1, assign.expand(*x.shape[:-1], -1))
means = torch.as_tensor(np.ascontiguousarray(result['means']),
device='cuda', dtype=dtype)
components = torch.as_tensor(np.ascontiguousarray(result['components']),
device='cuda', dtype=dtype)
comp2d = components.reshape(-1, result['D'])
g = x @ means.T
h2d = x @ comp2d.T
K, r = components.shape[0], components.shape[1]
coeffs = torch.as_tensor(np.ascontiguousarray(result['coeffs']),
device='cuda', dtype=dtype)
corr = torch.zeros(*x.shape[:-1], result['n'], device='cuda', dtype=dtype)
for k in range(K):
mask = assign == k
if mask.any():
ids = torch.nonzero(mask, as_tuple=False).flatten()
corr[..., ids] = h2d[..., k * r:(k + 1) * r] @ coeffs[ids].T
return g.gather(-1, assign.expand(*x.shape[:-1], -1)) + corr