"""第 38 章：真实训练、权重 INT8/INT4、QAT、剪枝与蒸馏；CPU 小实验。"""

import argparse
from copy import deepcopy
import hashlib
import json
from pathlib import Path
import platform
from tempfile import TemporaryDirectory
import torch
from torch import nn
from torch.nn import functional as F
from digits_project import make_data

CONFIG = dict(teacher_width=64, student_width=16, epochs=80, qat_epochs=20,
              batch_size=64, lr=.003, qat_lr=.0003, weight_decay=.0001,
              temperature=2., alpha=.7, teacher_seed=2038, student_seed=2039,
              order_seed=2040, max_accuracy_drop=.01, max_state_bytes=6000)


def quantize_rows(weight, bits):
    """对 Linear 的每个输出通道求一个 scale；使用对称范围，保留一个未用码。"""
    if bits not in (4, 8) or weight.ndim != 2 or not torch.isfinite(weight).all():
        raise ValueError("需要有限二维权重，位数为 4 或 8")
    limit = 2**(bits-1)-1
    maximum = weight.detach().abs().amax(dim=1, keepdim=True)
    scale = torch.where(maximum > 0, maximum/limit, torch.ones_like(maximum))
    codes = torch.round(weight.detach()/scale).clamp(-limit, limit).to(torch.int8)
    return codes, scale


def fake_weight(weight, bits):
    codes, scale = quantize_rows(weight, bits)
    rounded = codes.to(weight.dtype)*scale
    # 前向使用舍入后的值；反向对浮点主权重采用恒等直通估计，不对 scale 求导。
    return weight+(rounded-weight).detach()


def pack_int4(codes):
    if codes.dtype != torch.int8 or torch.any(codes < -7) or torch.any(codes > 7):
        raise ValueError("本例的对称 4 位码范围为 -7 到 7")
    unsigned = (codes.flatten().to(torch.int16)+7).to(torch.uint8)
    if unsigned.numel() % 2:
        unsigned = torch.cat([unsigned, torch.zeros(1, dtype=torch.uint8)])
    return unsigned[::2] | (unsigned[1::2] << 4)


def unpack_int4(packed, count):
    if packed.dtype != torch.uint8 or count < 1 or packed.numel() != (count+1)//2:
        raise ValueError("压缩字节数与原元素数不匹配")
    unsigned = torch.stack([packed & 15, packed >> 4], dim=1).flatten()[:count]
    if torch.any(unsigned > 14):
        raise ValueError("遇到本格式未使用的码 15")
    return (unsigned.to(torch.int16)-7).to(torch.int8)


class PackedLinear(nn.Module):
    def __init__(self, layer, bits):
        super().__init__()
        codes, scale = quantize_rows(layer.weight, bits)
        self.bits = bits
        self.rows, self.columns = layer.weight.shape
        self.register_buffer("codes", pack_int4(codes) if bits == 4 else codes)
        self.register_buffer("scale", scale.clone())
        self.register_buffer("bias", layer.bias.detach().clone())

    def weight(self):
        codes = (unpack_int4(self.codes, self.rows*self.columns).reshape(self.rows, self.columns)
                 if self.bits == 4 else self.codes)
        return codes.float()*self.scale

    def forward(self, x):
        # 真正紧凑保存，但每次解包回 FP32 计算；这里没有高效整数矩阵内核。
        return F.linear(x, self.weight(), self.bias)


class Classifier(nn.Module):
    def __init__(self, width, fake_bits=None):
        super().__init__()
        self.width, self.fake_bits = width, fake_bits
        self.first, self.last = nn.Linear(64, width), nn.Linear(width, 10)

    def forward(self, x):
        x = x.flatten(1)
        for index, layer in enumerate([self.first, self.last]):
            x = (F.linear(x, fake_weight(layer.weight, self.fake_bits), layer.bias)
                 if self.fake_bits is not None else layer(x))
            if index == 0:
                x = x.relu()
        return x


def pack_model(model, bits):
    result = deepcopy(model)
    result.fake_bits = None
    result.first, result.last = PackedLinear(model.first, bits), PackedLinear(model.last, bits)
    return result.eval()


def state_bytes(model):
    return sum(t.numel()*t.element_size() for t in model.state_dict().values())


def save_model(model, path):
    if model.fake_bits is not None:
        raise ValueError("请先把 QAT 模型转换为 PackedLinear 再导出")
    bits = model.first.bits if isinstance(model.first, PackedLinear) else None
    torch.save(dict(format=1, width=model.width, bits=bits,
                    state=model.state_dict()), path)


def load_model(path):
    data = torch.load(path, map_location="cpu", weights_only=True)
    if data["format"] != 1 or data["width"] not in (16, 64) or data["bits"] not in (None, 4, 8):
        raise ValueError("不支持的模型格式")
    model = Classifier(data["width"])
    if data["bits"] is not None:
        model = pack_model(model, data["bits"])
    model.load_state_dict(data["state"], strict=True)
    return model.eval()


def distillation_loss(student, teacher, labels, temperature=2., alpha=.7):
    hard = F.cross_entropy(student, labels)
    soft = F.kl_div(F.log_softmax(student/temperature, dim=1),
                   F.softmax(teacher.detach()/temperature, dim=1),
                   reduction="batchmean")
    return (1-alpha)*hard+alpha*temperature**2*soft


@torch.inference_mode()
def evaluate(model, dataset):
    model.eval()
    x, y = dataset
    logits = model(x)
    prediction = logits.argmax(1)
    return dict(loss=F.cross_entropy(logits, y).item(),
                accuracy=(prediction == y).float().mean().item(),
                predictions=prediction.tolist())


def fit(model, train, validation, teacher=None, epochs=None, lr=None):
    epochs = CONFIG["epochs"] if epochs is None else epochs
    lr = CONFIG["lr"] if lr is None else lr
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=CONFIG["weight_decay"])
    generator = torch.Generator().manual_seed(CONFIG["order_seed"])
    x, y = train
    teacher_logits = None
    if teacher is not None:
        teacher.eval()
        with torch.no_grad():
            teacher_logits = teacher(x).detach()
    best_state, best_loss, best_epoch = deepcopy(model.state_dict()), evaluate(model, validation)["loss"], 0
    history = []
    for epoch in range(1, epochs+1):
        model.train()
        for ids in torch.randperm(len(x), generator=generator).split(CONFIG["batch_size"]):
            optimizer.zero_grad(set_to_none=True)
            logits = model(x[ids])
            loss = (F.cross_entropy(logits, y[ids]) if teacher is None else
                    distillation_loss(logits, teacher_logits[ids], y[ids],
                                      CONFIG["temperature"], CONFIG["alpha"]))
            if not torch.isfinite(loss):
                raise RuntimeError("训练损失非有限")
            loss.backward()
            optimizer.step()
        score = evaluate(model, validation)
        history.append(dict(epoch=epoch, loss=score["loss"], accuracy=score["accuracy"]))
        if score["loss"] < best_loss:
            best_state, best_loss, best_epoch = deepcopy(model.state_dict()), score["loss"], epoch
    model.load_state_dict(best_state)
    return model.eval(), best_epoch, history


def prune_weights(model, fraction=.9):
    result = deepcopy(model)
    weights = [result.first.weight, result.last.weight]
    flat = torch.cat([w.detach().flatten() for w in weights])
    mask = torch.ones_like(flat)
    mask[flat.abs().argsort()[:int(fraction*len(flat))]] = 0
    start = 0
    with torch.no_grad():
        for weight in weights:
            count = weight.numel()
            weight.mul_(mask[start:start+count].reshape(weight.shape))
            start += count
    return result


def prune_units(model, train_x, keep=16):
    # 只用训练输入估计隐藏单元贡献，选择后整行/整列删除，不做额外微调。
    with torch.inference_mode():
        activation = model.first(train_x.flatten(1)).relu().mean(0)
        importance = activation*model.last.weight.abs().sum(0)
        chosen = importance.topk(keep).indices.sort().values
    result = Classifier(keep)
    with torch.no_grad():
        result.first.weight.copy_(model.first.weight[chosen])
        result.first.bias.copy_(model.first.bias[chosen])
        result.last.weight.copy_(model.last.weight[:, chosen])
        result.last.bias.copy_(model.last.bias)
    return result.eval()


def run():
    torch.set_num_threads(1)
    torch.use_deterministic_algorithms(True)
    raw, ids = make_data()
    datasets = [(x.float()/16, y) for x, y in raw]
    train, validation, test = datasets
    torch.manual_seed(CONFIG["teacher_seed"])
    teacher, teacher_epoch, teacher_history = fit(Classifier(64), train, validation)
    torch.manual_seed(CONFIG["student_seed"])
    initial = Classifier(16)
    hard, hard_epoch, hard_history = fit(deepcopy(initial), train, validation)
    student, student_epoch, student_history = fit(deepcopy(initial), train, validation, teacher)
    qat = deepcopy(teacher)
    qat.fake_bits = 4
    qat, qat_epoch, qat_history = fit(qat, train, validation, epochs=CONFIG["qat_epochs"],
                                     lr=CONFIG["qat_lr"])
    models = dict(teacher=teacher, ptq8=pack_model(teacher, 8), ptq4=pack_model(teacher, 4),
                  qat4=pack_model(qat, 4), prune90=prune_weights(teacher),
                  units16=prune_units(teacher, train[0]), hard16=hard, distill16=student)
    validation_scores = {name: evaluate(model, validation) for name, model in models.items()}
    eligible = [name for name, model in models.items()
                if name != "teacher" and state_bytes(model) <= CONFIG["max_state_bytes"]
                and validation_scores[name]["accuracy"] >=
                validation_scores["teacher"]["accuracy"]-CONFIG["max_accuracy_drop"]]
    selected = min(eligible, key=lambda name: (state_bytes(models[name]),
                                              validation_scores[name]["loss"])) if eligible else "teacher"
    # 到这里所有方案和选择均已确定；随后用已在旧章公开的测试划分做能力回归。
    reference = torch.tensor(evaluate(teacher, test)["predictions"])
    target = test[1]
    rows = []
    with TemporaryDirectory() as directory:
        for name, model in models.items():
            path = Path(directory)/"model.pt"
            save_model(model, path)
            restored = load_model(path)
            with torch.inference_mode():
                torch.testing.assert_close(model(test[0]), restored(test[0]), atol=0, rtol=0)
            score = evaluate(restored, test)
            prediction = torch.tensor(score["predictions"])
            regressed = (reference == target) & (prediction != target)
            fixed = (reference != target) & (prediction == target)
            weights = [model.first.weight(), model.last.weight()] if isinstance(model.first, PackedLinear) else [
                model.first.weight, model.last.weight]
            rows.append(dict(name=name, state_bytes=state_bytes(model),
                             weight_values=sum(w.numel() for w in weights),
                             nonzero_weights=sum(int((w != 0).sum()) for w in weights),
                             validation=validation_scores[name], test=score,
                             regressions=ids[2][regressed.numpy()].tolist(),
                             fixes=ids[2][fixed.numpy()].tolist(),
                             class_correct=[int(((prediction == target) & (target == digit)).sum())
                                            for digit in range(10)]))
    digest = hashlib.sha256()
    for x, y in datasets:
        digest.update(x.numpy().tobytes())
        digest.update(y.numpy().tobytes())
    report = dict(complete=True, config=CONFIG, selected=selected, data_sha256=digest.hexdigest(),
                  environment=dict(python=platform.python_version(), torch=torch.__version__,
                                   device="cpu", dtype="float32", threads=1),
                  split_ids=[part.tolist() for part in ids], test_targets=target.tolist(),
                  class_counts=[int((target == digit).sum()) for digit in range(10)],
                  epochs=dict(teacher=teacher_epoch, hard16=hard_epoch,
                              distill16=student_epoch, qat4=qat_epoch),
                  history=dict(teacher=teacher_history, hard16=hard_history,
                               distill16=student_history, qat4=qat_history), rows=rows)
    return report, models


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--out", type=Path)
    parser.add_argument("--save-dir", type=Path)
    args = parser.parse_args()
    report, models = run()
    if args.out:
        args.out.parent.mkdir(parents=True, exist_ok=True)
        args.out.write_text(json.dumps(report, ensure_ascii=False, indent=2)+"\n")
    if args.save_dir:
        args.save_dir.mkdir(parents=True, exist_ok=True)
        for name, model in models.items():
            save_model(model, args.save_dir/f"{name}.pt")
    print("沿用数字划分：训练/验证/回归测试 = 1077 / 360 / 360")
    print("按验证交叉熵选轮次：", report["epochs"])
    for row in report["rows"]:
        print(f"{row['name']}：状态={row['state_bytes']} 字节；"
              f"验证准确率={row['validation']['accuracy']:.3f}；"
              f"测试准确率={row['test']['accuracy']:.3f}；"
              f"退步/修正={len(row['regressions'])}/{len(row['fixes'])}")
    print("按预设验证条件选择：", report["selected"])
    print("全部模型保存/加载后的分数一致；INT4 每字节实际存两个码。")
    print("量化层解包回 FP32 运算，未测整数内核加速；状态字节不等于峰值内存。")


if __name__ == "__main__":
    main()

1. 0.88 除以 0.25 得 3.52，整数码为 4，还原成 1.0，绝对误差为 0.12。2.0 的整数码原本应为 8，截为 7 后还原成 1.75，误差为 0.25。
2. 每项仍占 8 位。打包后，一个字节保存两个 4 位码；元素总数为奇数时，最后一个字节只有一半装着有效数据。
3. 权重占 4736 字节，74 个 FP32 步长占 296 字节，74 个 FP32 偏置再占 296 字节，合计 5328 字节。
4. 张量的形状和数据类型没变，每个零仍占 4 字节。换用适合的稀疏格式并计入索引，或删除整行整列使形状变小，才可能减少存储。是否加速还要单独测。
5. 假量化改变的是前向使用的数值；主权重仍按 FP32 存储。本章必须转换为 PackedLinear 并打包，才得到紧凑状态。保存后还要加载验证。
6. 方向是 $D_{\mathrm{KL}}(p_t\|p_s)$。最高分类别保留下来，其余候选的相对概率及教师有多犹豫都丢失了。教师概率也可能有错，所以本例还使用真实标签。
7. 一张标签为 1 的图从正确变成错误，另一张标签为 5 的图从错误变成正确，正确总数相抵。按类别或样本检查，才能看到这个变化。
8. 每个种子下，两个学生仍应复制同一个初始模型、使用同样的批次顺序。用训练集更新、验证集比较；不要根据回归测试分数挑一个最好种子。若继续改方案，最终泛化结论需要新的独立评估。
9. 可先使用与 QAT 相同的额外轮数作为一组对照，但还要记录训练规模、每轮工作量及教师预训练成本。报告恢复训练前后的结果，也保留没有恢复到预期能力的情况。