This page is the module’s teaching text. It takes the NumPy fluency of Modules 1–2 and adds the one thing NumPy lacks — a gradient — then builds up the working vocabulary of PyTorch that the DSP-and-ML labs of Course 3 Module 8 assume: tensors and devices, autograd, nn.Module, the data pipeline, a training and evaluation loop written by hand, the signal and image transforms, quantization, and the export path from a trained model to onnxruntime on the Mac and TensorRT on the Jetson. The arc follows the UvA notebook tutorial Introduction to PyTorch and PyTorch’s own Learn the Basics; both remain available as optional deep-dives, and nothing below requires them. As with every lessons page in this course, it is AI-drafted teaching text reviewed by me; see the syllabus’s note on AI use.
Two framing rules for the module. First, the training loop is written by hand — Lightning, Trainer-style wrappers, and Hugging Face’s accelerate are popular and fine in production, but the loop is short, and owning it is what makes the deployment questions (what exactly runs on the board?) answerable. Second, the model’s job ends in a file: a state_dict, an ONNX graph, and a saved reference output that the C and Rust modules — and Course 3’s edge labs — are verified against with np.testing.assert_allclose. PyTorch is the arbiter, not the product.
1 · Tensors: ndarrays with a device and a gradient
A torch.Tensor is an n-dimensional array with the NumPy model of Module 1 underneath — shape, dtype, strides, views vs. copies, the same broadcasting rules — plus two attributes NumPy does not have: a device (where the memory lives) and, optionally, a gradient.
1.1 Creation and interop
import numpy as np, torchx = torch.zeros(2, 3) # float32 by default — not float64 like NumPyr = torch.rand(2, 3, 4) # uniform [0, 1)a = torch.arange(6).reshape(2, 3) # int64t = torch.tensor([[1.0, 2.0], [3.0, 4.0]])np_arr = np.linspace(0, 1, 8, dtype=np.float32)v = torch.from_numpy(np_arr) # zero-copy: shares memory with np_arrv[0] =99.0# ...so np_arr[0] is now 99.0 tooback = v.numpy() # zero-copy in the other direction (CPU tensors only)
Three defaults differ from NumPy and cause most first-week bugs:
NumPy
PyTorch
Consequence
Default float dtype
float64
float32
torch.from_numpy of a float64 array gives a float64 tensor; feeding it to a float32 model raises a dtype mismatch. Convert at the boundary: .astype(np.float32) or .float().
Default int dtype
int64
int64
Same — but labels for CrossEntropyLoss must be int64 (torch.long), not int32.
torch.Tensor(2, 3)
—
uninitialized memory
The uppercase constructor allocates without filling; use torch.zeros/torch.empty deliberately.
torch.from_numpy and .numpy() share memory, which is the property the Course 3 host-in-the-loop labs rely on: a capture decoded with np.frombuffer becomes a tensor for free, and a model output becomes a NumPy array for assert_allclose for free. The sharing stops at a device boundary — a tensor on the GPU must come back with .cpu() before .numpy().
1.2 Devices
Every tensor and every model lives on a device. The code is written once against a device variable and the hardware is selected at the top of the script:
torch.accelerator is the portable form of the older "cuda" if torch.cuda.is_available() else "cpu" idiom and covers Apple’s MPS backend on the Mac as well as CUDA on the Linux box (Course 4 Lab 1.1’s machine) and the Jetson. Two rules follow: everything that touches a tensor in the training loop must be on the same device (inputs, labels, model), and anything going to NumPy, Matplotlib, or a file comes back through .cpu(). MPS is a first-class training backend for small models but has gaps — some operators fall back to the CPU silently or raise; float64 is unsupported — so the course’s convention is develop on MPS, verify on CPU (Exercise 3.1 records which operators fall back).
1.3 Shapes, views, and the layout convention
view, reshape, permute, transpose, squeeze/unsqueeze, flatten, cat/stack are the NumPy operations under new names; the strides model of Module 1 §2 applies unchanged. view requires a contiguous tensor and is always zero-copy; reshape copies when it must; permute returns a non-contiguous view, and .contiguous() materializes it.
The batch-first layout convention is the one thing to internalize: nn layers expect (N, C, L) for 1-D signals (Conv1d), (N, C, H, W) for images (Conv2d), and (N, L, C) or (L, N, C) for recurrent layers depending on batch_first. A mono audio frame of 16 000 samples is (1, 1, 16000); a log-mel patch of 40 bands × 100 frames is (1, 1, 40, 100); a batch of 32 of them is (32, 1, 40, 100). Most shape errors are a missing channel dimension — x.unsqueeze(1) — and the error message names the expected rank.
1.4 Arithmetic, broadcasting, and in-place operations
Broadcasting is NumPy’s (Module 1 §4); @/torch.matmul batches over leading dimensions exactly as np.matmul does; torch.einsum takes the same subscript language as Module 1 §4.2, and the same reading — kept indices name the output, repeated-and-absent indices are summed — carries over unchanged. The family maps one-to-one: torch.tensordot(a, b, dims=…), torch.outer, torch.kron, torch.trace/torch.diagonal, and torch.bmm for the strictly-3-D batched product @ generalizes. Two things are PyTorch-specific. torch.einsum dispatches to the same batched-GEMM kernels as matmul whenever the string is a (batched) matrix product, so on the GPU the two spellings cost the same; and with the opt_einsum package installed, torch.backends.opt_einsum makes three-or-more-operand strings choose a contraction order automatically (Module 1 §4.2’s optimize= concern, handled for you — check torch.backends.opt_einsum.is_available()). The idiom that appears most is the attention-style score, written both ways so the equivalence is visible:
S = torch.einsum("bqd,bkd->bqk", Q, K) # index formS = Q @ K.transpose(-1, -2) # matmul form — identical kernels
Every einsum here is differentiable like any other op, so the index string is also the cleanest way to write a custom layer’s forward pass before profiling tells you whether a fused matmul spelling is worth the loss of clarity. Operations with a trailing underscore (x.add_(1), x.mul_(0.5), x.zero_()) mutate in place — useful for memory, forbidden on tensors that autograd still needs (the error message says so). Reductions take dim= where NumPy takes axis=, and keepdim= where NumPy takes keepdims=.
2 · Autograd
2.1 Define-by-run
PyTorch records a computation graph as the operations run — a define-by-run framework — and backward() walks it in reverse to fill each leaf’s .grad. The UvA tutorial’s example is the whole mechanism in seven lines:
x = torch.arange(3, dtype=torch.float32, requires_grad=True)a = x +2b = a **2c = b +3y = c.mean() # scalary.backward() # dy/dx for every leaf with requires_gradprint(x.grad) # tensor([1.3333, 2.0000, 2.6667]) — check it by hand: 2(x+2)/3
The rules that matter in practice:
Rule
Why it bites
backward()accumulates into .grad
Two calls without optimizer.zero_grad() (or x.grad = None) sum gradients from two iterations — the classic silent training bug.
Only leaf tensors with requires_grad=True get .grad
Intermediate results (a, b, c above) do not keep a gradient unless .retain_grad() is called.
The graph is freed after backward()
A second backward() on the same graph raises; retain_graph=True exists for the rare legitimate case.
Evaluation and inference run inside one of these — otherwise every forward pass builds a graph that is never used and memory grows.
.detach() cuts a tensor out of the graph; .item() extracts a Python scalar
Logging loss instead of loss.item() keeps the whole graph alive through the list of losses.
In-place ops on a needed tensor raise
Autograd checks a version counter; the error names the operation.
This is the Course 1 Lesson 7 gradient made mechanical: every torch function knows its own local derivative, and the chain rule is applied by the graph walk. Exercise 3.2 verifies backward() against a hand derivative and against a finite-difference check — the same discipline as checking a NumPy kernel against SciPy.
2.2 What autograd is not
It is not symbolic differentiation (no closed forms), not forward-mode by default (torch.func.jvp exists), and not free: every intermediate that the backward pass needs is kept alive until backward() runs, which is why memory scales with depth × batch size and why no_grad matters for inference. torch.func.grad and vmap provide a functional interface for the cases where a gradient of a gradient, or a per-sample gradient, is wanted.
2.3 Gradient descent by hand, then the optimizer
Before torch.optim hides it, one step of gradient descent on a least-squares fit is three lines — and they are exactly what optimizer.step() and optimizer.zero_grad() do:
w = torch.zeros(2, requires_grad=True) # the parametersfor _ inrange(steps): loss = ((X @ w - y) **2).mean() # Course 1 Lesson 6's objective loss.backward() # w.grad = 2 Xᵀ(Xw − y) / Nwith torch.no_grad(): # the update is not part of the graph w -= lr * w.grad # optimizer.step() w.grad =None# optimizer.zero_grad()
The no_grad block is the point: updating a parameter is an in-place operation on a leaf that autograd would otherwise record. torch.optim.SGD generalizes the two lines to every parameter, adds momentum and weight decay, and Adam adds the per-parameter scaling (Course 1 Lesson 42); optimizer.state_dict() is where that extra state lives, which is why a resumable checkpoint saves it.
3 · nn.Module and the layer catalog
3.1 Anatomy of a module
import torch.nn as nnclass SimpleClassifier(nn.Module):def__init__(self, num_inputs: int, num_hidden: int, num_outputs: int):super().__init__()self.linear1 = nn.Linear(num_inputs, num_hidden)self.act_fn = nn.Tanh()self.linear2 = nn.Linear(num_hidden, num_outputs)def forward(self, x): x =self.linear1(x) x =self.act_fn(x)returnself.linear2(x) # logits — no sigmoid here (see §4)
Layers assigned as attributes in __init__ are registered: their parameters appear in model.parameters() (what the optimizer updates), in model.state_dict() (what is saved), and they move with model.to(device). A layer stored in a plain Python list is not registered — use nn.ModuleList or nn.Sequential. forward is called through model(x), never directly, because the call operator runs hooks and mode checks.
model = SimpleClassifier(2, 4, 1)n_params =sum(p.numel() for p in model.parameters()) # count before you deployfor name, p in model.named_parameters():print(name, tuple(p.shape)) # linear1.weight (4, 2) ...
Counting parameters before training is a Course 3 habit: Lab 8.2’s edge budget is stated in parameters, and a state_dict is what eventually becomes a C array.
Out of scope for the labs; named so the catalog is complete
nn.Sequential(...) chains layers when forward is just a pipeline; a custom forward is needed as soon as there is a skip connection, a second input, or a reshape between blocks. A conv block for a 1-D signal is the recurring unit:
Initialization is automatic (Kaiming-uniform for Linear/Conv); torch.nn.init overrides it when a paper or a numerical problem demands.
3.3 Tracing shapes through a network
Every shape in a conv net is hand-derivable, and deriving them before running is how a layout bug is caught at the design stage rather than in a stack trace. For a batch of 32 log-mel patches with 40 bands and 100 frames, through three conv_blocks (kernel 5, padding=2, then MaxPool 2) and an adaptive-pool head:
Stage
Shape (N, C, H, W)
Rule
Input
(32, 1, 40, 100)
one channel, bands × frames
conv_block(1, 8)
(32, 8, 20, 50)
padding = k // 2 keeps H, W; pool halves both
conv_block(8, 16)
(32, 16, 10, 25)
same
conv_block(16, 32)
(32, 32, 5, 12)
pooling floors: 25 → 12
AdaptiveAvgPool2d(1)
(32, 32, 1, 1)
any spatial size → 1 × 1
Flatten
(32, 32)
Linear(32, K)
(32, K)
logits, one per class
The general formula for a convolution or pool along one axis is \(L_{\text{out}} = \lfloor (L_{\text{in}} + 2p - d(k-1) - 1)/s \rfloor + 1\) with padding \(p\), dilation \(d\), kernel \(k\), stride \(s\). The adaptive pool is what makes the head independent of the number of frames — and what a fixed-shape export (§9) later pins back down. Parameter count follows the same table: a Conv2d(C_in, C_out, k) holds \(C_{\text{in}} C_{\text{out}} k^2 + C_{\text{out}}\) numbers, a Linear(a, b) holds \(ab + b\); BatchNorm adds \(2C\) trainable plus \(2C\) running statistics. Summing these by hand for the table above is Exercise 3.4’s budget line.
3.4 Dtypes that deploy
dtype
Where it is used
Note
float32
Training, the reference
The arbiter’s precision
float16 / bfloat16
Mixed-precision training (§10); fp16 inference on the Jetson
bfloat16 keeps float32’s exponent range with fewer mantissa bits — safer for training, absent on most edge runtimes
int8 / uint8 (qint8/quint8 in the quantization APIs)
4 · Losses, optimizers, and the numerics behind them
4.1 Losses
Loss
Expects
Notes
nn.CrossEntropyLoss()
logits (N, K), labels (N,) of int64
Combines log_softmax and negative log-likelihood — the cross-entropy of Course 1 Lesson 43, computed stably. Do not apply softmax before it.
nn.BCEWithLogitsLoss()
logits (N,) or (N, 1), targets float in {0, 1}
Sigmoid folded in, for the same stability reason; the UvA XOR classifier uses it.
nn.MSELoss(), nn.L1Loss(), nn.HuberLoss()
prediction and target of equal shape
Regression and denoising (Lab 8.3’s spectral mask is trained with an L1 or MSE on the masked magnitude)
nn.NLLLoss()
log-probabilities
Only when log_softmax is already part of the model
The “logits in, no activation at the end of the model” convention is not style: log(softmax(x)) computed as two steps overflows for large logits, and the fused losses use the log-sum-exp trick. It also means the deployed model emits logits, and the argmax or threshold lives in the application code — a detail that matters when the application code is C on the Jetson.
SGD with momentum, Adam, and AdamW are the Course 1 Lesson 42 gradient algorithms with per-parameter step-size adaptation; weight_decay is the ℓ₂ penalty of Lesson 40 applied inside the step (decoupled in AdamW). A scheduler changes the learning rate on a fixed program (StepLR, CosineAnnealingLR, OneCycleLR) or on a plateau (ReduceLROnPlateau); scheduler.step() is called once per epoch (or per batch for OneCycleLR — read the docstring). torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) between backward() and step() bounds an exploding gradient, which recurrent models need.
TipWorking rule
Start with Adam(lr=1e-3), no scheduler, and a model small enough to overfit a single batch. If it cannot memorize eight samples, the bug is in the shapes, the labels, or the loss — not in the hyperparameters.
A Dataset answers “how many” and “give me item i”; a DataLoader batches, shuffles, and optionally parallelizes (num_workers) and pins memory for faster host→GPU copies (pin_memory=True, CUDA only). Batching stacks items along a new leading dimension with the default collate_fn; variable-length audio needs a custom one (pad to the longest, return lengths). drop_last=True keeps every batch the same shape, which matters once BatchNorm is involved or a fixed-shape export is planned.
For the labs’ data — recorded keywords, synthetic vibration signatures, noisy/clean audio pairs — the Dataset is where the Module 1–2 preprocessing lives: soundfile.read → scipy.signal.resample_poly → librosa.feature.melspectrogram → torch.from_numpy(...).float(). Precompute features to .npy once and index them, rather than recomputing an STFT in every epoch.
5.2 Transforms
torchvision.transforms.v2 (Compose, ToImage, ToDtype, Normalize, RandomCrop, …) is the image pipeline; torchaudio.transforms (Resample, Spectrogram, MelSpectrogram, MFCC, AmplitudeToDB) is the audio one — and, unlike the librosa versions, these run inside the model on the device and export with it (§7). Augmentation for signals is cheap and effective: random time shifts, additive noise at a sampled SNR, SpecAugment-style frequency masks (torchaudio.transforms.FrequencyMasking).
6 · The training loop, written by hand
6.1 The loop
def train_epoch(model, loader, loss_fn, optimizer, device): model.train() # Dropout on, BatchNorm updating total =0.0for x, y in loader: x, y = x.to(device), y.to(device) logits = model(x).squeeze(dim=1) # (N, 1) -> (N,) for BCEWithLogits loss = loss_fn(logits, y.float()) optimizer.zero_grad() loss.backward() optimizer.step() total += loss.item() * x.shape[0]return total /len(loader.dataset)@torch.no_grad()def evaluate(model, loader, device): model.eval() # Dropout off, BatchNorm frozen correct =0for x, y in loader: x, y = x.to(device), y.to(device) pred = (torch.sigmoid(model(x).squeeze(dim=1)) >=0.5).long() correct += (pred == y).sum().item()return correct /len(loader.dataset)
Five steps per batch — move, forward, loss, backward, step — and two mode switches that are easy to forget and silent when forgotten: model.train() before training, model.eval() before evaluation. The outer loop calls train_epoch, then evaluate on a held-out split, logs both, and keeps the best checkpoint:
Early stopping is the loop above with a patience counter. A confusion matrix (sklearn.metrics.confusion_matrix on the collected predictions, Module 2) is the evaluation artifact the labs ask for, and per-class accuracy is what reveals a keyword the model never learned.
Save the state_dict, not the module object — the object pickles Python class paths that break on refactor, and weights_only=True refuses arbitrary pickle payloads. A checkpoint for resuming training additionally holds the optimizer state, the scheduler state, and the epoch, in one dict.
Reproducibility has three layers: seed everything (torch.manual_seed, np.random.default_rng(seed), the DataLoader’s generator=), request determinism (torch.use_deterministic_algorithms(True), torch.backends.cudnn.benchmark = False), and accept that some GPU kernels are non-deterministic regardless — the notes record the seed and the observed run-to-run spread rather than pretending it is zero.
6.3 Watching it train
torch.utils.tensorboard.SummaryWriter writes scalars, histograms, and figures to a log directory; tensorboard --logdir runs serves them. The three plots worth having on every run are train loss per epoch, validation metric per epoch, and the learning rate — a validation curve that rises then falls is overfitting, one that never moves is a data or label bug. add_figure takes the Matplotlib figures of Module 1, so the confusion matrix and a few misclassified spectrograms go in the same log.
6.4 Metrics beyond accuracy
Accuracy hides a class the model never predicts. Collect the logits and labels for the whole validation set once (a list of .cpu() tensors, concatenated at the end) and hand them to Module 2’s sklearn.metrics: confusion_matrix, classification_report (precision, recall, F1 per class), and for a binary detector the ROC curve and AUC from the raw logits — the detection-theory picture of Course 3 Lab 6.7 theory, where the threshold is a design choice rather than 0.5 by default. For the denoiser the metric is a signal metric, not a classification one: SNR improvement in dB computed in NumPy from the clean, noisy, and output waveforms, exactly as Course 3 Lab 8.3 reports it. Keep metric computation in NumPy/scikit-learn on the CPU; the model’s only job is the forward pass.
6.5 What the wrappers do
pytorch-lightning, torch.compile-based trainers, and accelerate wrap exactly the loop above: the mode switches, device moves, logging, checkpointing, and mixed precision. They are worth adopting once the loop is understood, and they are worth avoiding until then — every deployment question (“what does the forward pass receive?”, “is the normalization inside the model?”) is answered by reading the loop you wrote.
7 · Signals and images inside the graph
7.1 torch.fft and torchaudio
torch.fft.rfft, irfft, rfftfreq, stft/istft mirror np.fft/scipy.signal (Module 1) and are differentiable, so a spectral loss — an L1 on log-magnitude — is one line. torchaudio.transforms.MelSpectrogram(sample_rate, n_fft, hop_length, n_mels) followed by AmplitudeToDB reproduces the librosa log-mel of Module 2 §1 as a module; put it inside the model and the deployed graph takes raw audio, which removes a whole class of “the featurizer on the board does not match the one in training” bugs (Lab 8.2’s Analysis section is about exactly that). The two libraries disagree in defaults (center, window, mel scale htk vs. slaney, power vs. amplitude); Exercise 3.4 pins them against each other with assert_allclose before training starts.
7.2 torchvision and transfer learning
torchvision.models ships pretrained classifiers (resnet18, mobilenet_v3_small) and detectors (fasterrcnn_*, ssdlite320_mobilenet_v3_large, retinanet_*) with a weights= enum whose .transforms() gives the matching preprocessing. Transfer learning replaces the head and fine-tunes: freeze the backbone (requires_grad_(False)), swap model.fc/model.classifier[-1] for a Linear with the new class count, train the head, then optionally unfreeze the last block with a smaller learning rate. Lab 8.6’s COCO detector is the pretrained path with no training at all — export, verify, deploy — and its reference boxes come from running the torchvision model on the Mac in eval() mode, which is the arbiter the TensorRT engine is checked against.
8 · Making it small: quantization and pruning
A trained float32 model is a starting point, not a deployable. The edge ladder is fp32 → fp16 → int8, and each rung is a Course 1 Lesson 45 rate–distortion trade made concrete: fewer bits per weight and activation against a measured accuracy loss.
Mode
What is quantized
When the scales are chosen
Cost / benefit
fp16 / bf16
Everything, by casting
Never — a format change
Halves memory and doubles GPU throughput; trtexec --fp16 needs no calibration (Lab 8.6)
Dynamic int8
Weights statically; activations per-batch at run time
At inference
Trivial to apply; helps Linear/RNN-heavy models; little for convs
Static PTQ (post-training)
Weights and activations
From a calibration set run through the model once
The edge default; needs representative data — Lab 8.2’s calibration patches
QAT (quantization-aware training)
Same as static
Simulated during fine-tuning (“fake quant”)
Recovers most of the accuracy static PTQ loses; costs a training run
Two API facts to hold loosely, because they are moving. In PyTorch’s torch.ao.quantization package the eager-mode functions (quantize_dynamic, prepare/convert with a QConfig, fuse_modules, QuantStub/DeQuantStub) and the FX functions (prepare_fx/convert_fx) are in maintenance, and the documentation states that quantization development is being centralized in the torchao package — the PT2-export workflow (prepare_pt2e/convert_pt2e with a backend quantizer such as the XNNPACK one, which targets ARM CPUs) lives there. Check which functions the installed versions expose before writing Exercise 3.6, and record the versions in the notes.
The second fact is that for this site’s targets, quantization mostly happens after export, in the deployment tool: onnxruntime.quantization.quantize_static/quantize_dynamic produce an int8 ONNX; TensorRT’s int8 path calibrates on the Jetson with an IInt8Calibrator (Lab 8.2 has the details and the trtexec caveats). The PyTorch-side quantization matters most for QAT, where the fake-quant graph is what makes the exported int8 model accurate. torch.nn.utils.prune implements magnitude pruning (structured and unstructured); on a CPU or GPU without sparse kernels it saves nothing at run time unless the pruned channels are physically removed, which is why the labs measure size and latency separately.
ImportantThe arbiter rule
The float32 model on the Mac produces the reference outputs. Every smaller or faster variant — fp16, int8, the ONNX graph, the TensorRT engine, the C port — is compared against those outputs with np.testing.assert_allclose at a stated tolerance, and the tolerance is part of the deliverable. “Looks right” is not a verification.
9 · Export and edge deployment
9.1 From nn.Module to ONNX
model.eval()example = torch.randn(1, 1, 40, 100) # one log-mel patchprog = torch.onnx.export(model, (example,), dynamo=True, # torch.export-based exporter input_names=["logmel"], output_names=["logits"], dynamic_shapes={"x": {0: "batch"}}) # or omit for a fixed shapeprog.save("kws.onnx")import onnxonnx.checker.check_model(onnx.load("kws.onnx"))
Since PyTorch 2.9 the torch.export-based exporter (dynamo=True) is the default: it traces the model into a graph with Python control flow removed, and returns an ONNXProgram whose .save() writes the file. The older TorchScript-based exporter (dynamo=False, dynamic_axes=) still exists for models the new path cannot trace; torch.onnx.export(model, args, "file.onnx", input_names=..., output_names=...) is the legacy one-liner most tutorials still show. Choose a fixed input shape for the edge (Lab 8.2 does), because both TensorRT and the microcontroller ports want static buffers; pick the opset_version the target runtime supports and record it. Export aftermodel.eval() so BatchNorm uses running statistics and Dropout is gone, and confirm with onnx.checker.
9.2 Running the ONNX
import onnxruntime as ortsess = ort.InferenceSession("kws.onnx", providers=["CPUExecutionProvider"])name = sess.get_inputs()[0].namey_onnx = sess.run(None, {name: x_np})[0] # NumPy in, NumPy outy_ref = model(torch.from_numpy(x_np)).detach().numpy()np.testing.assert_allclose(y_onnx, y_ref, rtol=1e-4, atol=1e-5)np.savez("kws_reference.npz", x=x_np, logits=y_ref) # the arbiter file
onnxruntime runs the same file on the Mac (CPUExecutionProvider, or CoreMLExecutionProvider when installed), on the Pi 5 (CPU, NEON via the default provider), and on the Jetson (CUDAExecutionProvider, TensorrtExecutionProvider when the onnxruntime-gpu Jetson wheel is installed). The parity check against the PyTorch output is the first thing run on every new machine, and the saved .npz is what Course 3’s on-device scripts load.
9.3 Reading the exported graph
The ONNX file is a protobuf of nodes, initializers (the weights), and typed inputs and outputs; onnx.helper.printable_graph(model.graph) prints it, and the Netron desktop app draws it. Three things to look for on the first export of any model:
BatchNorm is gone. The exporter folds it into the preceding convolution’s weights and bias in eval() mode. If a BatchNormalization node survives, the model was exported in training mode.
The featurizer’s fate.torch.stft maps to ONNX’s STFT operator only from opset 17, and runtimes support it unevenly — TensorRT in particular. The usual edge decision is the one Lab 8.2 makes: export the classifier alone with a log-mel input, and reimplement the featurizer on the board in C, verified against the Python one with the same .npz. Exercise 3.4 builds the model so that split is one attribute access.
Node count and opset. The dynamo=True exporter tends to produce a flatter graph than the legacy tracer; the count is a cheap regression check, and the opset must be one the target runtime lists as supported (onnxruntime and TensorRT publish their ranges).
9.4 TensorRT on the Jetson, and the rest of the zoo
flowchart LR T["nn.Module<br/>float32, Mac"] --> S["state_dict .pt"] T --> O["ONNX<br/><i>source of truth</i>"] O --> R1["onnxruntime<br/>Mac · Pi 5 · Jetson"] O --> E["trtexec --fp16 / --int8<br/>.engine, <b>per-board artifact</b>"] O --> L["LiteRT / tflite int8<br/>Pi 5 (Lab 8.2)"] T --> C["weights as .npy → C / Rust port<br/>(Modules 4–12, Lesson 38)"] T --> A["reference .npz"] A -. assert_allclose .-> R1 & E & L & C
flowchart LR
T["nn.Module<br/>float32, Mac"] --> S["state_dict .pt"]
T --> O["ONNX<br/><i>source of truth</i>"]
O --> R1["onnxruntime<br/>Mac · Pi 5 · Jetson"]
O --> E["trtexec --fp16 / --int8<br/>.engine, <b>per-board artifact</b>"]
O --> L["LiteRT / tflite int8<br/>Pi 5 (Lab 8.2)"]
T --> C["weights as .npy → C / Rust port<br/>(Modules 4–12, Lesson 38)"]
T --> A["reference .npz"]
A -. assert_allclose .-> R1 & E & L & C
TensorRT ships with JetPack (/usr/src/tensorrt/bin/trtexec); an engine is built on the board from the ONNX, is specific to that JetPack/TensorRT/GPU combination, and is never copied between machines — the ONNX is the artifact under version control. LiteRT (the renamed TensorFlow Lite) is the other common edge runtime; the ONNX→tflite conversion tools exist and Lab 8.2 uses one for the Pi’s int8 path. Apple’s Core ML (coremltools) is the equivalent on the Mac and iOS and is out of scope here. Below all of these sits the hand port: a state_dict exported as .npy arrays becomes const float tables in C or a static slice in Rust, and a Conv1d is the FIR bank the C modules already know how to write — CMSIS-NN provides the int8 kernels on the Cortex-M, and that is where the ML labs’ “does it fit on the STM32?” question is eventually answered.
Measure before tuning.torch.utils.benchmark.Timer times a statement with warm-up and proper synchronization (torch.cuda.synchronize()/torch.mps.synchronize() are implied — the naive time.perf_counter() around a GPU op measures the launch, not the work, exactly Course 4 Lab 1.1’s lesson). torch.profiler.profile(activities=[...], record_shapes=True) around a few training steps, then prof.key_averages().table(sort_by="self_cpu_time_total"), shows whether the time is in the model, the DataLoader, or host↔︎device copies — the usual answer for small models is the loader.
Mixed precision.torch.autocast(device_type="cuda", dtype=torch.float16) around the forward and loss, plus torch.amp.GradScaler for the backward, halves memory and speeds up tensor-core GPUs; on MPS autocast supports float16 with fewer guarantees. Keep the loss and the optimizer state in float32 — the Course 1 Lesson 37 conditioning argument applies to gradient accumulation too.
torch.compile(model) fuses and specializes the graph for a real speed-up on CUDA after a warm-up compile; on MPS its coverage is partial. Leave it off while debugging; turn it on when timing.
DataLoader knobs.num_workers > 0 moves decoding and augmentation off the training process (on macOS the workers spawn, so the script needs the if __name__ == "__main__": guard); persistent_workers=True avoids re-forking every epoch.
11 · What transfers to C and Rust
PyTorch concept
On the Cortex-M / in the C or Rust port
Where it is done
state_dict
const weight tables (.npy → header or Rust static)
Module 4/5 data layout, Module 6 storage
nn.Linear
Dot products; dot_q15 from Module 0 with a bias
Modules 4–5 kernels
nn.Conv1d
An FIR bank; CMSIS-DSP/CMSIS-NN on the M4
Module 8 (talking to hardware), Course 3 Module 6
BatchNorm
Folded into the conv weights at export — nothing at run time
The host-side test harness the C/Rust kernel is checked with
Module 12 testing; Course 3 Module 6 harness
The port never re-derives anything: the Python model defines the numbers, the .npz defines the expected outputs, and the C or Rust implementation is correct when the arbiter says so at the stated tolerance — the same three-step discipline Module 2 §5 set up for a filter, applied to a network.
12 · Lesson → exercise map
Section
Exercise it feeds
§1 tensors, interop, devices, layout
3.1 (interop and device probe)
§2 autograd
3.2 (by hand vs. backward)
§3 nn.Module, §4 losses/optimizers, §5 data, §6 the loop
3.3 (XOR from scratch), 3.4 (keyword CNN), 3.5 (denoiser)