PyTorch
实战

from pathlib import Path
import platform
import sys

def project_report():
    root = Path(__file__).resolve().parents[1]
    return {
        "python": sys.version.split()[0],
        "platform": platform.platform(),
        "executable": sys.executable,
        "project_root": str(root),
        "expected_dirs": [
            name for name in ("src", "scripts", "configs", "runs", "tests")
            if (root / name).exists()
        ],
    }

for key, value in project_report().items():
    print(f"{key}: {value}")

assert sys.version_info >= (3, 10), "需要 Python 3.10 或更高版本"

import random
import platform
import sys

import torch

def seed_everything(seed):
    random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

def report(seed=7):
    seed_everything(seed)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    sample = torch.randn(2, 3, device=device)
    return {
        "python": sys.version.split()[0],
        "platform": platform.platform(),
        "torch": torch.__version__,
        "cuda_build": torch.version.cuda,
        "cuda_available": torch.cuda.is_available(),
        "device": str(device),
        "device_name": torch.cuda.get_device_name(0) if device.type == "cuda" else "CPU",
        "sample_mean": round(sample.mean().item(), 6),
    }

for key, value in report().items():
    print(f"{key}: {value}")

import json
from pathlib import Path

import torch
from torch import nn

def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    torch.manual_seed(7)
    model = nn.Linear(4, 3).to(device)
    features = torch.tensor([[1.0, 0.0, 0.5, -1.0], [0.2, 0.7, 0.1, 0.3]], device=device)
    targets = torch.tensor([1, 2], dtype=torch.long, device=device)
    logits = model(features)
    loss = nn.CrossEntropyLoss()(logits, targets)
    loss.backward()
    grad_norm = torch.linalg.vector_norm(model.weight.grad).item()
    assert logits.shape == (2, 3)
    assert torch.isfinite(loss) and torch.isfinite(model.weight.grad).all()
    report = {"device": str(device), "loss": loss.item(), "grad_norm": grad_norm}
    Path("runs").mkdir(exist_ok=True)
    Path("runs/first-run.json").write_text(
        json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    print(report)

if __name__ == "__main__":
    main()

import json
from pathlib import Path
import sys

import torch

def make_manifest():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    return {
        "python": sys.version.split()[0],
        "torch": torch.__version__,
        "cuda_build": torch.version.cuda,
        "device": str(device),
        "project": "pytorch-practice",
        "schema": 1,
    }

manifest = make_manifest()
Path("runs").mkdir(exist_ok=True)
Path("runs/environment.json").write_text(
    json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8"
)
assert json.loads(Path("runs/environment.json").read_text(encoding="utf-8")) == manifest
print(manifest)

import torch

scores = torch.arange(12, dtype=torch.float32).reshape(3, 4)
print("scores:", scores)
print("first row:", scores[0])
print("second column:", scores[:, 1])
print("middle block:", scores[:, 1:3])

row_keep = scores[1:2]
row_drop = scores[1]
print("keep axis:", row_keep.shape, "drop axis:", row_drop.shape)
assert row_keep.shape == (1, 4)
assert row_drop.shape == (4,)

mask = scores.ge(6)
selected = scores[mask]
print("selected shape:", selected.shape)
assert selected.tolist() == [6.0, 7.0, 8.0, 9.0, 10.0, 11.0]

copied = scores.clone()
copied[0, 0] = -1
assert scores[0, 0].item() == 0.0

import torch

batch = torch.arange(12, dtype=torch.float32).reshape(3, 4)
offset = torch.tensor([[10.0, 20.0, 30.0, 40.0]])
shifted = batch + offset
print("broadcast:", shifted.shape)
assert shifted.shape == (3, 4)

weights = torch.randn(4, 2)
hidden = batch @ weights
print("matmul:", hidden.shape)
assert hidden.shape == (3, 2)

mean_features = batch.mean(dim=0)
mean_samples = batch.mean(dim=1, keepdim=True)
print("feature mean:", mean_features.shape)
print("sample mean:", mean_samples.shape)
assert mean_features.shape == (4,)
assert mean_samples.shape == (3, 1)

parts = torch.split(batch, 2, dim=1)
joined = torch.cat(parts, dim=1)
assert torch.equal(joined, batch)

import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
features = torch.randn(2, 3, dtype=torch.float32, device=device)
labels = torch.tensor([0, 1], dtype=torch.long, device=device)
mask = labels.eq(1)
print(features.dtype, features.device, features.is_contiguous())
print(labels.dtype, labels.device, mask.dtype)

matrix = torch.arange(24, dtype=torch.float32).reshape(2, 3, 4)
permuted = matrix.permute(0, 2, 1)
print("strides:", matrix.stride(), permuted.stride())
contiguous = permuted.contiguous()
assert contiguous.is_contiguous()
assert contiguous.shape == (2, 4, 3)

try:
    bad = features.to("cpu") + labels.to(device)
except RuntimeError as error:
    print("device mismatch caught:", error.__class__.__name__)

import torch

def check_features(features, batch, feature_count):
    assert features.shape == (batch, feature_count)
    assert features.dtype == torch.float32
    assert torch.isfinite(features).all()

def check_labels(labels, batch, classes):
    assert labels.shape == (batch,)
    assert labels.dtype == torch.long
    assert int(labels.min()) >= 0 and int(labels.max()) < classes

features = torch.tensor([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]], dtype=torch.float32)
labels = torch.tensor([1, 0], dtype=torch.long)
check_features(features, batch=2, feature_count=3)
check_labels(labels, batch=2, classes=2)

indices = torch.tensor([[2, 0], [1, 1]], dtype=torch.long)
table = torch.arange(12).reshape(3, 4)
selected = table.gather(0, indices)
print("selected:", selected)
assert selected.shape == indices.shape

from dataclasses import dataclass

import torch
from torch.utils.data import Dataset

@dataclass(frozen=True)
class Record:
    sample_id: str
    features: tuple[float, ...]
    label: int

class TinyDataset(Dataset):
    def __init__(self, records, feature_count):
        self.records = list(records)
        self.feature_count = feature_count
        if not self.records:
            raise ValueError("数据集不能为空")
        labels = [record.label for record in self.records]
        if min(labels) < 0:
            raise ValueError("标签必须非负")

    def __len__(self):
        return len(self.records)

    def __getitem__(self, index):
        record = self.records[index]
        if len(record.features) != self.feature_count:
            raise ValueError(f"{record.sample_id}: 特征数错误")
        features = torch.tensor(record.features, dtype=torch.float32)
        label = torch.tensor(record.label, dtype=torch.long)
        return {"features": features, "label": label, "id": record.sample_id}

records = [Record("a", (0.1, 0.2), 0), Record("b", (0.8, 0.4), 1)]
dataset = TinyDataset(records, feature_count=2)
print(len(dataset), dataset[0])
assert dataset[0]["features"].shape == (2,)

import torch
from torch.utils.data import DataLoader, TensorDataset

features = torch.arange(20, dtype=torch.float32).reshape(10, 2)
labels = torch.arange(10, dtype=torch.long) % 2
dataset = TensorDataset(features, labels)
loader = DataLoader(dataset, batch_size=3, shuffle=False, drop_last=False)

seen = []
for batch_index, (batch_features, batch_labels) in enumerate(loader):
    print(batch_index, batch_features.shape, batch_labels.tolist())
    assert batch_features.ndim == 2
    assert batch_features.shape[0] == batch_labels.shape[0]
    seen.extend(batch_labels.tolist())

assert len(seen) == len(dataset)
assert len(loader) == 4
print("all labels seen:", seen)

from collections import Counter

def audit_split(train_ids, valid_ids, test_ids, labels):
    sets = {
        "train": set(train_ids),
        "valid": set(valid_ids),
        "test": set(test_ids),
    }
    overlaps = {
        "train_valid": sets["train"] & sets["valid"],
        "train_test": sets["train"] & sets["test"],
        "valid_test": sets["valid"] & sets["test"],
    }
    if any(overlaps.values()):
        raise ValueError(f"发现 ID 交集: {overlaps}")
    label_summary = {name: Counter(labels[name]) for name in labels}
    return {
        "sizes": {name: len(values) for name, values in sets.items()},
        "labels": {name: dict(counts) for name, counts in label_summary.items()},
    }

report = audit_split(
    ["a", "b", "c"], ["d"], ["e"],
    {"train": [0, 1, 0], "valid": [1], "test": [0]},
)
print(report)

import random

import torch
from torch.utils.data import DataLoader, TensorDataset

def seed_worker(worker_id):
    seed = torch.initial_seed() % 2**32
    random.seed(seed)

generator = torch.Generator().manual_seed(7)
features = torch.arange(24, dtype=torch.float32).reshape(12, 2)
labels = torch.arange(12, dtype=torch.long)
dataset = TensorDataset(features, labels)
loader = DataLoader(
    dataset, batch_size=4, shuffle=True, num_workers=0,
    worker_init_fn=seed_worker, generator=generator,
)
first_epoch = [int(label) for _, batch_labels in loader for label in batch_labels]
generator = torch.Generator().manual_seed(7)
loader_again = DataLoader(dataset, batch_size=4, shuffle=True, num_workers=0, generator=generator)
second_epoch = [int(label) for _, batch_labels in loader_again for label in batch_labels]
assert first_epoch == second_epoch
assert sorted(first_epoch) == list(range(12))
print("deterministic order:", first_epoch)

import torch
from torch import nn

class TemperatureHead(nn.Module):
    def __init__(self, features, classes):
        super().__init__()
        self.classifier = nn.Linear(features, classes)
        self.temperature = nn.Parameter(torch.ones(()))
        self.register_buffer("running_scale", torch.tensor(1.0))

    def forward(self, features):
        logits = self.classifier(features)
        return logits / self.temperature.clamp_min(1e-4) * self.running_scale

model = TemperatureHead(4, 3)
print("parameters:", [name for name, _ in model.named_parameters()])
print("buffers:", [name for name, _ in model.named_buffers()])
assert "classifier.weight" in model.state_dict()
assert "temperature" in model.state_dict()
assert "running_scale" in model.state_dict()

import torch
from torch import nn

class ResidualMLP(nn.Module):
    def __init__(self, features, hidden, classes):
        super().__init__()
        self.input_layer = nn.Linear(features, hidden)
        self.block = nn.Sequential(nn.ReLU(), nn.Linear(hidden, hidden))
        self.output_layer = nn.Linear(hidden, classes)

    def forward(self, features, return_hidden=False):
        if features.ndim != 2 or features.shape[-1] != self.input_layer.in_features:
            raise ValueError("输入必须是 (B, features)")
        hidden = self.input_layer(features)
        hidden = hidden + self.block(hidden)
        logits = self.output_layer(torch.relu(hidden))
        if return_hidden:
            return {"logits": logits, "hidden": hidden}
        return logits

model = ResidualMLP(4, 8, 3)
result = model(torch.randn(5, 4), return_hidden=True)
print(result["hidden"].shape, result["logits"].shape)
assert result["logits"].shape == (5, 3)

import torch
from torch import nn

model = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 3))
records = []

def make_hook(name):
    def hook(module, inputs, output):
        tensor = output if isinstance(output, torch.Tensor) else output[0]
        records.append({
            "name": name,
            "shape": tuple(tensor.shape),
            "dtype": str(tensor.dtype),
            "device": str(tensor.device),
            "mean": tensor.detach().float().mean().item(),
            "finite": bool(torch.isfinite(tensor).all()),
        })
    return hook

handles = [module.register_forward_hook(make_hook(name)) for name, module in model.named_modules() if name]
logits = model(torch.randn(6, 4))
for handle in handles:
    handle.remove()
print(records)
assert records[-1]["shape"] == (6, 3)
assert torch.isfinite(logits).all()

import torch

def train_one_epoch(model, loader, optimizer, criterion, device):
    model.train()
    total_loss = 0.0
    total_items = 0
    for features, targets in loader:
        features = features.to(device)
        targets = targets.to(device)
        optimizer.zero_grad(set_to_none=True)
        logits = model(features)
        loss = criterion(logits, targets)
        if loss.ndim != 0:
            raise ValueError("训练 loss 必须是标量")
        loss.backward()
        optimizer.step()
        count = targets.shape[0]
        total_loss += loss.detach().item() * count
        total_items += count
    return total_loss / max(total_items, 1)

@torch.no_grad()
def evaluate(model, loader, criterion, device):
    model.eval()
    total_loss = 0.0
    total_items = 0
    total_correct = 0
    for features, targets in loader:
        logits = model(features.to(device))
        targets = targets.to(device)
        loss = criterion(logits, targets)
        total_loss += loss.item() * targets.shape[0]
        total_items += targets.shape[0]
        total_correct += (logits.argmax(dim=1) == targets).sum().item()
    return {"loss": total_loss / max(total_items, 1), "accuracy": total_correct / max(total_items, 1)}

from dataclasses import dataclass, asdict
import json

@dataclass
class EpochLog:
    epoch: int
    samples_seen: int
    train_loss: float
    valid_loss: float
    accuracy: float
    learning_rate: float
    seconds: float

history = []
history.append(EpochLog(1, 128, 0.83, 0.79, 0.71, 1e-3, 2.4))
history.append(EpochLog(2, 256, 0.55, 0.61, 0.80, 1e-3, 2.3))
with open("runs/history.json", "w", encoding="utf-8") as handle:
    json.dump([asdict(row) for row in history], handle, ensure_ascii=False, indent=2)
best = min(history, key=lambda row: row.valid_loss)
print("best epoch:", best.epoch, "samples:", best.samples_seen)

from pathlib import Path
import torch

def save_checkpoint(path, model, optimizer, epoch, best_metric, config):
    payload = {
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "epoch": epoch,
        "best_metric": best_metric,
        "config": config,
        "torch_version": torch.__version__,
    }
    Path(path).parent.mkdir(parents=True, exist_ok=True)
    torch.save(payload, path)

def load_checkpoint(path, model, optimizer=None, device="cpu"):
    payload = torch.load(path, map_location=device, weights_only=False)
    model.load_state_dict(payload["model"])
    if optimizer is not None and "optimizer" in payload:
        optimizer.load_state_dict(payload["optimizer"])
    return payload["epoch"], payload["best_metric"], payload["config"]

import torch

def run_step(model, batch, optimizer, criterion, device, update=True):
    features, targets = batch
    features = features.to(device)
    targets = targets.to(device)
    if update:
        model.train()
        optimizer.zero_grad(set_to_none=True)
    else:
        model.eval()
    with torch.set_grad_enabled(update):
        logits = model(features)
        loss = criterion(logits, targets)
        if update:
            loss.backward()
            optimizer.step()
    return {"loss": float(loss.detach()), "count": int(targets.shape[0]), "mode": "train" if update else "eval"}

from pathlib import Path

import torch
from torchvision import datasets, transforms

root = Path("data")
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,)),
])
train_set = datasets.MNIST(root, train=True, download=True, transform=transform)
test_set = datasets.MNIST(root, train=False, download=True, transform=transform)

image, label = train_set[0]
print("train/test:", len(train_set), len(test_set))
print("image:", image.shape, image.dtype, image.min().item(), image.max().item())
print("label:", label, type(label))
assert image.shape == (1, 28, 28)
assert image.dtype == torch.float32
assert 0 <= int(label) <= 9

import torch
from torch import nn

class MnistCNN(nn.Module):
    def __init__(self, classes=10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.classifier = nn.Linear(32 * 7 * 7, classes)

    def forward(self, images):
        if images.ndim != 4 or images.shape[1:] != (1, 28, 28):
            raise ValueError(f"期望 (B,1,28,28)，得到 {tuple(images.shape)}")
        features = self.features(images)
        return self.classifier(features.flatten(1))

model = MnistCNN()
batch = torch.randn(4, 1, 28, 28)
logits = model(batch)
print(logits.shape)
assert logits.shape == (4, 10)
assert torch.isfinite(logits).all()

import torch

def confusion_matrix(targets, predictions, classes=10):
    matrix = torch.zeros(classes, classes, dtype=torch.long)
    for truth, prediction in zip(targets.tolist(), predictions.tolist()):
        matrix[truth, prediction] += 1
    return matrix

def collect_errors(model, loader, device, limit=20):
    model.eval()
    errors = []
    with torch.no_grad():
        for images, targets in loader:
            logits = model(images.to(device))
            probabilities = logits.softmax(dim=1)
            predictions = probabilities.argmax(dim=1).cpu()
            for index, (truth, prediction) in enumerate(zip(targets, predictions)):
                if int(truth) != int(prediction) and len(errors) < limit:
                    errors.append({
                        "target": int(truth),
                        "prediction": int(prediction),
                        "confidence": float(probabilities[index, prediction]),
                    })
    return errors

# matrix = confusion_matrix(targets, predictions)
# print(matrix)

import torch
from torch import nn

class LinearBaseline(nn.Module):
    def __init__(self, classes=10):
        super().__init__()
        self.net = nn.Sequential(nn.Flatten(), nn.Linear(28 * 28, classes))

    def forward(self, images):
        return self.net(images)

class MLPBaseline(nn.Module):
    def __init__(self, hidden=128, classes=10):
        super().__init__()
        self.net = nn.Sequential(
            nn.Flatten(), nn.Linear(28 * 28, hidden), nn.ReLU(),
            nn.Linear(hidden, classes),
        )

    def forward(self, images):
        return self.net(images)

def probe(model):
    images = torch.randn(4, 1, 28, 28)
    logits = model(images)
    return tuple(logits.shape), sum(p.numel() for p in model.parameters())

for model in (LinearBaseline(), MLPBaseline()):
    print(type(model).__name__, probe(model))

import torch
from torch import nn

def check_batch(features, targets, feature_count, classes):
    if features.ndim != 2 or features.shape[1] != feature_count:
        raise ValueError(f"features shape: {tuple(features.shape)}")
    if targets.ndim != 1 or targets.shape[0] != features.shape[0]:
        raise ValueError("targets 必须与 batch 对齐")
    if targets.dtype != torch.long:
        raise TypeError(f"targets dtype: {targets.dtype}")
    if targets.min().item() < 0 or targets.max().item() >= classes:
        raise ValueError("targets 超出类别范围")

model = nn.Linear(4, 3)
features = torch.tensor([[0.1, 0.2, 0.3, 0.4]])
targets = torch.tensor([1], dtype=torch.long)
check_batch(features, targets, feature_count=4, classes=3)
logits = model(features)
loss = nn.CrossEntropyLoss()(logits, targets)
assert torch.isfinite(loss)
print("minimal reproduction passed")

import torch
from torch import nn

def test_classifier_contract():
    model = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 3))
    features = torch.randn(5, 4)
    labels = torch.tensor([0, 1, 2, 1, 0], dtype=torch.long)
    logits = model(features)
    assert logits.shape == (5, 3)
    assert logits.dtype == features.dtype
    assert torch.isfinite(logits).all()
    loss = nn.CrossEntropyLoss()(logits, labels)
    assert loss.ndim == 0 and torch.isfinite(loss)
    before = [parameter.detach().clone() for parameter in model.parameters()]
    loss.backward()
    assert all(parameter.grad is not None for parameter in model.parameters())
    changed = any(not torch.equal(old, new.detach()) for old, new in zip(before, model.parameters()))
    assert not changed, "backward 不应自动改变参数"

test_classifier_contract()
print("contract tests passed")

import torch
from torch import nn

model = nn.Sequential(nn.Linear(256, 512), nn.ReLU(), nn.Linear(512, 10))
optimizer = torch.optim.SGD(model.parameters(), lr=1e-2)
inputs = torch.randn(32, 256)
targets = torch.randint(0, 10, (32,))
criterion = nn.CrossEntropyLoss()

with torch.profiler.profile(
    activities=[torch.profiler.ProfilerActivity.CPU],
    record_shapes=True,
    profile_memory=True,
) as profile:
    for _ in range(3):
        optimizer.zero_grad(set_to_none=True)
        loss = criterion(model(inputs), targets)
        loss.backward()
        optimizer.step()

print(profile.key_averages().table(sort_by="cpu_time_total", row_limit=5))

import json
from pathlib import Path
import torch

def summarize_batch(features, targets, batch_index, sample_ids):
    return {
        "batch_index": batch_index,
        "shape": tuple(features.shape),
        "dtype": str(features.dtype),
        "device": str(features.device),
        "finite": bool(torch.isfinite(features).all()),
        "target_shape": tuple(targets.shape),
        "target_min": int(targets.min()),
        "target_max": int(targets.max()),
        "sample_ids": list(sample_ids),
    }

# snapshot = summarize_batch(features, targets, 0, ids)
# Path("runs").mkdir(exist_ok=True)
# Path("runs/debug-batch.json").write_text(json.dumps(snapshot, indent=2), encoding="utf-8")

from pathlib import Path

import torch
from torch import nn

class Classifier(nn.Module):
    def __init__(self, features=4, classes=3):
        super().__init__()
        self.linear = nn.Linear(features, classes)

    def forward(self, features):
        return self.linear(features)

model = Classifier()
sample = torch.tensor([[0.1, 0.2, 0.3, 0.4]])
reference = model(sample).detach()
path = Path("runs/classifier.pt")
path.parent.mkdir(parents=True, exist_ok=True)
torch.save({"model": model.state_dict(), "config": {"features": 4, "classes": 3}}, path)

restored = Classifier(features=4, classes=3)
payload = torch.load(path, map_location="cpu", weights_only=False)
restored.load_state_dict(payload["model"], strict=True)
restored_output = restored(sample).detach()
torch.testing.assert_close(reference, restored_output)
print("round trip passed:", path)

import torch
from torch import nn

class Predictor:
    def __init__(self, model, device="cpu", labels=None):
        self.model = model.to(device)
        self.model.eval()
        self.device = torch.device(device)
        self.labels = labels or [str(index) for index in range(10)]

    @torch.inference_mode()
    def __call__(self, images):
        tensor = torch.as_tensor(images, dtype=torch.float32)
        if tensor.ndim == 2:
            tensor = tensor.unsqueeze(0).unsqueeze(0)
        elif tensor.ndim == 3:
            tensor = tensor.unsqueeze(0)
        if tensor.ndim != 4 or tensor.shape[1] != 1:
            raise ValueError(f"期望 (B,1,H,W)，实际 {tuple(tensor.shape)}")
        logits = self.model(tensor.to(self.device))
        probabilities = logits.softmax(dim=-1)
        indices = probabilities.argmax(dim=-1).cpu().tolist()
        return [{"label": self.labels[index], "index": index, "probability": float(probabilities[row, index])}
                for row, index in enumerate(indices)]

# predictor = Predictor(model, labels=[str(i) for i in range(10)])
# print(predictor(torch.zeros(28, 28)))

import json
from pathlib import Path

required = {"model", "config", "preprocess", "labels", "data_version", "code_version"}
manifest = {
    "model": "mnist-cnn-v1",
    "config": "configs/mnist-cnn.json",
    "preprocess": "normalize-0.1307-0.3081",
    "labels": "labels-v1.json",
    "data_version": "mnist-v1",
    "code_version": "commit-or-date",
}
missing = required - manifest.keys()
if missing:
    raise ValueError(f"manifest 缺少: {sorted(missing)}")
Path("runs/release-manifest.json").write_text(
    json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8"
)
loaded = json.loads(Path("runs/release-manifest.json").read_text(encoding="utf-8"))
assert loaded == manifest
print("release manifest passed")

from dataclasses import dataclass, asdict
import json
from pathlib import Path

@dataclass(frozen=True)
class BundleManifest:
    model_version: str
    preprocess_version: str
    labels: tuple[str, ...]
    input_shape: tuple[int, int, int]
    code_version: str

manifest = BundleManifest(
    model_version="mnist-cnn-v1",
    preprocess_version="normalize-v1",
    labels=tuple(str(index) for index in range(10)),
    input_shape=(1, 28, 28),
    code_version="2026-08-28",
)
bundle = Path("release/mnist-v1")
bundle.mkdir(parents=True, exist_ok=True)
(bundle / "manifest.json").write_text(
    json.dumps(asdict(manifest), ensure_ascii=False, indent=2), encoding="utf-8"
)
loaded = json.loads((bundle / "manifest.json").read_text(encoding="utf-8"))
assert tuple(loaded["input_shape"]) == manifest.input_shape
assert len(loaded["labels"]) == 10
print("bundle ready:", bundle)

from pathlib import Path
import json

import torch
from torch import nn
from torch.utils.data import DataLoader, random_split
from torchvision import datasets, transforms

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
        )
        self.head = nn.Linear(32 * 7 * 7, 10)

    def forward(self, images):
        return self.head(self.features(images).flatten(1))

def evaluate(model, loader, device):
    model.eval()
    correct = total = 0
    with torch.no_grad():
        for images, labels in loader:
            logits = model(images.to(device))
            labels = labels.to(device)
            correct += (logits.argmax(1) == labels).sum().item()
            total += labels.numel()
    return correct / max(total, 1)

def main(epochs=2, batch_size=128):
    torch.manual_seed(7)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,)),
    ])
    full = datasets.MNIST("data", train=True, download=True, transform=transform)
    train_set, valid_set = random_split(full, [54000, 6000], generator=torch.Generator().manual_seed(7))
    test_set = datasets.MNIST("data", train=False, download=True, transform=transform)
    train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True)
    valid_loader = DataLoader(valid_set, batch_size=batch_size)
    test_loader = DataLoader(test_set, batch_size=batch_size)
    model = Net().to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)
    criterion = nn.CrossEntropyLoss()
    history = []
    best = 0.0
    for epoch in range(1, epochs + 1):
        model.train()
        for images, labels in train_loader:
            optimizer.zero_grad(set_to_none=True)
            loss = criterion(model(images.to(device)), labels.to(device))
            loss.backward()
            optimizer.step()
        valid_accuracy = evaluate(model, valid_loader, device)
        history.append({"epoch": epoch, "valid_accuracy": valid_accuracy})
        if valid_accuracy > best:
            best = valid_accuracy
            Path("runs").mkdir(exist_ok=True)
            torch.save({"model": model.state_dict(), "epoch": epoch}, "runs/mnist-best.pt")
    print({"best_valid": best, "test": evaluate(model, test_loader, device)})
    Path("runs/history.json").write_text(json.dumps(history, indent=2), encoding="utf-8")

if __name__ == "__main__":
    main()

from dataclasses import dataclass, asdict
import json

@dataclass(frozen=True)
class RunSpec:
    run_id: str
    model: str
    weight_decay: float
    augmentation: str
    batch_size: int
    seed: int

runs = [
    RunSpec("mlp-s7", "mlp", 0.0, "none", 64, 7),
    RunSpec("cnn-s7", "cnn", 0.0, "none", 64, 7),
    RunSpec("cnn-decay-s7", "cnn", 1e-4, "none", 64, 7),
    RunSpec("cnn-aug-s7", "cnn", 1e-4, "shift", 64, 7),
]
with open("runs/ablation.json", "w", encoding="utf-8") as handle:
    json.dump([asdict(run) for run in runs], handle, ensure_ascii=False, indent=2)
for run in runs:
    assert run.batch_size > 0
    assert run.model in {"mlp", "cnn"}
print("planned runs:", len(runs))

from pathlib import Path
import json

import torch

def check_release(predictor, sample, expected_index, manifest_path):
    manifest = json.loads(Path(manifest_path).read_text(encoding="utf-8"))
    result = predictor(sample)[0]
    if result["index"] != expected_index:
        raise AssertionError(f"预测 {result['index']} != 期望 {expected_index}")
    if not 0.0 <= result["probability"] <= 1.0:
        raise AssertionError("概率超出范围")
    required = {"model", "preprocess", "labels", "code_version"}
    missing = required - manifest.keys()
    if missing:
        raise ValueError(f"manifest 缺少 {sorted(missing)}")
    return {"index": result["index"], "version": manifest["model"]}

# result = check_release(predictor, sample, expected_index=7, manifest_path="release/manifest.json")
# print(result)