Source browser
scripts/train_lenet.py
Read the complete file below or download its original bytes. Line numbers provide stable links to a declaration or processing step.
File fingerprint
1#!/usr/bin/env python3
2"""Train the documented modern LeNet variant and export plain portable weights.
3
4PyTorch is an optional preparation tool. The C++ examples load only float arrays.
5Training uses the official MNIST training split and evaluates the separate test split.
6"""
7import argparse # Describe the preparation inputs and reproducible training settings.
8import gzip # Read the original compressed IDX input files.
9import hashlib # Record exact training inputs and exported model bytes.
10from pathlib import Path # Resolve cache and output paths with pathlib.
11import struct # Decode the big-endian IDX headers from the published dataset.
12import time # Record completed training time for the preparation report.
13import data_sources # Share pinned MNIST inputs and the maintained Python provenance module.
14
15
16def digest(path): # Identify a file by its complete SHA-256 digest.
17 return hashlib.sha256(path.read_bytes()).hexdigest() # Hash the actual bytes used here.
18
19
20def read_idx(path, dimensions): # Validate and decode one official MNIST IDX archive.
21 raw = gzip.decompress(path.read_bytes()) # Expand only the named dataset archive.
22 magic = struct.unpack_from(">I", raw)[0] # Read the IDX type and dimension tag.
23 if magic != 0x0800 + dimensions: # Accept only unsigned-byte images or labels.
24 raise ValueError(f"unexpected IDX header in {path.name}") # Reject another data format.
25 shape = struct.unpack_from(">" + "I" * dimensions, raw, 4) # Read every dimension.
26 offset = 4 + 4 * dimensions # Locate the first payload byte after the header.
27 size = 1 # Begin the expected payload-size product.
28 for dimension in shape: # Include every declared array dimension.
29 size *= dimension # Count the unsigned-byte values in the archive.
30 if len(raw) != offset + size: # Check both truncation and trailing unexpected content.
31 raise ValueError(f"IDX payload size differs in {path.name}") # Report unsafe input.
32 return raw[offset:], shape # Return verified payload bytes and their shape.
33
34
35def main(): # Keep optional training dependencies outside ordinary C++ builds.
36 parser = argparse.ArgumentParser(description=__doc__) # Explain the preparation command.
37 parser.add_argument("--cache", type=Path, required=True) # Locate downloaded official IDX files.
38 parser.add_argument("--output", type=Path, default=Path("examples/data")) # Choose fixture storage.
39 parser.add_argument("--epochs", type=int, default=3) # Use a small but complete training run.
40 parser.add_argument("--batch-size", type=int, default=128) # Control each optimizer batch.
41 parser.add_argument("--seed", type=int, default=12345) # Reproduce initialization and batch order.
42 parser.add_argument("--threads", type=int, default=8) # Bound preparation CPU parallelism.
43 parser.add_argument("--metadata", type=Path, default=Path(data_sources.__file__).resolve()) # Store training results beside the maintained dataset records.
44 args = parser.parse_args() # Read the explicit training configuration.
45 if min(args.epochs, args.batch_size, args.threads) <= 0: # Reject empty or invalid training runs.
46 parser.error("epochs, batch size and thread count must be positive") # Explain the constraint.
47 import numpy as np # Export endian-defined arrays and decode verified dataset bytes.
48 import torch # Use the optional reference framework only during preparation.
49 from torch import nn # Construct the documented standard convolutional model.
50 torch.set_num_threads(args.threads) # Avoid using every host core for a small example.
51 torch.manual_seed(args.seed) # Fix both initial parameters and training permutation.
52 torch.use_deterministic_algorithms(True) # Select deterministic CPU preparation operations.
53
54 class LeNet(nn.Module): # Name the modern LeNet variant implemented by the native shaders.
55 def __init__(self): # Declare layers in exported parameter order.
56 super().__init__() # Initialize the framework's ordinary layer bookkeeping.
57 self.conv1 = nn.Conv2d(1, 6, 5) # Map padded input to six 28 by 28 feature maps.
58 self.conv2 = nn.Conv2d(6, 16, 5) # Map pooled features to sixteen 10 by 10 maps.
59 self.dense1 = nn.Linear(400, 120) # Apply the full C5-equivalent dense layer.
60 self.dense2 = nn.Linear(120, 84) # Apply the second hidden classifier layer.
61 self.dense3 = nn.Linear(84, 10) # Produce one logit per handwritten digit class.
62
63 def forward(self, values): # Match the native seven-dispatch inference pipeline.
64 values = nn.functional.pad(values, (2, 2, 2, 2)) # Add a zero border to 32 by 32.
65 values = nn.functional.avg_pool2d(torch.relu(self.conv1(values)), 2) # First conv and pool.
66 values = nn.functional.avg_pool2d(torch.relu(self.conv2(values)), 2) # Second conv and pool.
67 values = torch.relu(self.dense1(values.flatten(1))) # Flatten channel-major pooled maps.
68 values = torch.relu(self.dense2(values)) # Apply the second hidden ReLU layer.
69 return self.dense3(values) # Keep raw logits for stable cross entropy and comparison.
70
71 names = [ # Preserve the exact official split filenames in provenance.
72 "train-images-idx3-ubyte.gz", # Read all 60,000 training images.
73 "train-labels-idx1-ubyte.gz", # Read only their training labels during optimization.
74 "t10k-images-idx3-ubyte.gz", # Keep the 10,000-image test split held out.
75 "t10k-labels-idx1-ubyte.gz", # Use test labels only for final evaluation.
76 ] # Keep the training split independent of evaluation samples.
77 paths = [args.cache / name for name in names] # Find each verified archive in the public cache.
78 pinned = {entry["name"]: entry for entry in data_sources.SOURCE_RECORDS} # Reuse the exact download pins used by preparation.
79 for path in paths: # Check every original image and label archive before training reads it.
80 record = pinned[path.name] # Find this official split's expected compressed bytes.
81 if path.stat().st_size != record["bytes"] or digest(path) != record["sha256"]: # Reject changed inputs even when their IDX shape is plausible.
82 raise ValueError("MNIST source fingerprint differs: " + path.name) # Protect the recorded training and test provenance.
83 train_bytes, train_shape = read_idx(paths[0], 3) # Decode training images with validated shape.
84 train_labels, label_shape = read_idx(paths[1], 1) # Decode training class labels separately.
85 test_bytes, test_shape = read_idx(paths[2], 3) # Decode the held-out image split.
86 test_labels, test_label_shape = read_idx(paths[3], 1) # Decode the held-out class labels.
87 if train_shape != (60000, 28, 28) or test_shape != (10000, 28, 28): # Require official dimensions.
88 raise ValueError("the expected official MNIST split sizes are required") # Reject reshaped inputs.
89 if label_shape != (60000,) or test_label_shape != (10000,): # Pair every image with its own label.
90 raise ValueError("MNIST image and label counts differ") # Reject incomplete training inputs.
91 images = torch.from_numpy(np.frombuffer(train_bytes, dtype=np.uint8).copy()) # Own training bytes.
92 images = images.reshape(60000, 1, 28, 28).float().div_(255.0) # Match native zero-to-one scaling.
93 labels = torch.from_numpy(np.frombuffer(train_labels, dtype=np.uint8).copy()).long() # Class IDs.
94 test = torch.from_numpy(np.frombuffer(test_bytes, dtype=np.uint8).copy()) # Own test image bytes.
95 test = test.reshape(10000, 1, 28, 28).float().div_(255.0) # Apply identical test preprocessing.
96 targets = torch.from_numpy(np.frombuffer(test_labels, dtype=np.uint8).copy()).long() # Test classes.
97 model = LeNet() # Create the same layer shapes and parameter order as the C++ helper.
98 optimizer = torch.optim.Adam(model.parameters(), lr=0.001) # Learn all model weights and biases.
99 start = time.perf_counter() # Start timing after data preparation.
100 history = [] # Record each complete training epoch's mean cross entropy.
101 for epoch in range(args.epochs): # Visit the complete training split once per epoch.
102 model.train() # Make the preparation intent explicit for future layer changes.
103 order = torch.randperm(images.shape[0]) # Shuffle only the training examples.
104 loss_sum = 0.0 # Accumulate sample-weighted training loss for this epoch.
105 for offset in range(0, len(order), args.batch_size): # Include the final partial batch.
106 indices = order[offset:offset + args.batch_size] # Select this batch's original samples.
107 optimizer.zero_grad(set_to_none=True) # Release the previous batch's gradient values.
108 logits = model(images[indices]) # Evaluate the complete deep network.
109 loss = nn.functional.cross_entropy(logits, labels[indices]) # Compute supervised loss.
110 loss.backward() # Differentiate every layer in the reference preparation framework.
111 optimizer.step() # Update all trainable weights from the current batch.
112 loss_sum += float(loss.detach()) * len(indices) # Weight partial batches correctly.
113 history.append(loss_sum / len(order)) # Record completed sample-mean epoch loss.
114 print(f"epoch {epoch + 1}: training cross entropy {history[-1]:.6f}", flush=True) # Show progress.
115 training_seconds = time.perf_counter() - start # Record completed training duration.
116 model.eval() # Freeze preparation behavior before held-out evaluation.
117 with torch.no_grad(): # Preserve model parameters during evaluation.
118 logits = torch.cat([model(batch) for batch in test.split(256)]) # Bound evaluation memory.
119 accuracy = float((logits.argmax(1) == targets).float().mean()) # Measure all 10,000 labels.
120 test_loss = float(nn.functional.cross_entropy(logits, targets)) # Record held-out loss too.
121 weights = np.concatenate([p.detach().numpy().reshape(-1) for p in model.parameters()]) # Export order.
122 if weights.size != 61706: # Catch an accidental architecture or parameter-layout change.
123 raise ValueError("LeNet parameter layout must contain exactly 61,706 floats") # Protect loaders.
124 args.output.mkdir(parents=True, exist_ok=True) # Create only the requested fixture destination.
125 weight_path = args.output / "lenet_weights.bin" # Use a plain nonexecutable model representation.
126 weight_path.write_bytes(weights.astype("<f4").tobytes()) # Store portable little-endian float32.
127 reference_path = args.output / "lenet_reference.bin" # Keep framework logits as an independent oracle.
128 reference_path.write_bytes(logits.numpy().astype("<f4").tobytes()) # Save all 10,000 test predictions.
129 report = { # Save enough information to reproduce and audit the preparation run.
130 "architecture": "modern LeNet with full C3, ReLU, fixed average pooling and linear logits", # Name changes.
131 "parameters": int(weights.size), # Record the exact native buffer length.
132 "training_split": 60000, # State how many training samples informed weights.
133 "test_split": 10000, # State the independent evaluation population.
134 "seed": args.seed, # Record initialization and permutation seed.
135 "epochs": args.epochs, # Record complete training passes.
136 "batch_size": args.batch_size, # Record optimizer batch size.
137 "threads": args.threads, # Record preparation CPU parallelism.
138 "optimizer": "Adam", # Describe the reference preparation update rule.
139 "learning_rate": 0.001, # Record its initial learning rate.
140 "input_scale": "byte / 255 with a two-pixel zero border", # Describe exact preprocessing.
141 "training_cross_entropy": history, # Keep per-epoch learning progress.
142 "test_accuracy": accuracy, # Keep the exact measured held-out accuracy.
143 "test_cross_entropy": test_loss, # Keep a second classification-quality measure.
144 "training_seconds": training_seconds, # Keep optional preparation timing.
145 "torch_version": torch.__version__, # Identify the optional training implementation.
146 "numpy_version": np.__version__, # Identify the array export implementation.
147 "source_sha256": {path.name: digest(path) for path in paths}, # Identify all dataset inputs.
148 "weights_sha256": digest(weight_path), # Identify the actual deployed model bytes.
149 "reference_sha256": digest(reference_path), # Identify the independent test predictions.
150 } # Record the optional model-preparation environment.
151 artifacts = [dict(entry) for entry in data_sources.PREPARED_ARTIFACTS] # Preserve the prepared input records while updating the model export.
152 model_paths = {weight_path.name: weight_path, reference_path.name: reference_path} # Locate the two newly generated binary model artifacts.
153 for entry in artifacts: # Update model fingerprints and retain source attribution and byte layouts.
154 if entry["name"] in model_paths: # Leave every image, waveform, event and label fixture unchanged.
155 path = model_paths[entry["name"]] # Read only the newly exported model or framework reference.
156 entry["bytes"] = path.stat().st_size # Record the exact current binary export size.
157 entry["sha256"] = digest(path) # Bind the current model metadata to its actual exported bytes.
158 data_sources.write_metadata(artifacts, report, args.metadata) # Save current training provenance in the Python metadata module.
159 print(f"held-out MNIST accuracy: {accuracy:.4%}") # Report the complete test-split result.
160
161
162if __name__ == "__main__": # Run preparation only when this script is invoked explicitly.
163 main() # Train and export the requested portable example model.