Amanda: Unified Instrumentation Framework for Deep Neural Networks

Analyzing or optimizing a neural network — tracing its neurons, profiling its cost, pruning it, quantizing it — all means reaching into a running model, but every framework offers its own fragmented, model-specific hooks, so each tool is rewritten by hand for each model and backend; this work from Shanghai Jiao Tong, Microsoft, and others borrows the idea of instrumentation from program analysis and builds Amanda, which lets you write a tool once against a single before-and-after-operator abstraction and run it unchanged on any model in PyTorch or TensorFlow, cutting the code up to six times with under 5 percent overhead.

The argument in one glance

Problem

Every way of studying or shrinking a network means reaching into its execution

Neural networks are powerful but hard to deploy — expensive to compute, and opaque black boxes. So an enormous amount of research analyzes them (tracing which neurons fire, profiling their compute cost, extracting interpretable substructures) and optimizes them (pruning weights, quantizing tensors). Every one of these tasks shares a need: to monitor and manipulate the execution of an existing model — to reach in at the right place, read the right tensors, and sometimes change them. But the frameworks that run models, PyTorch and TensorFlow, offer only limited, ad-hoc support for doing this.
What people did

Use the framework's built-in hooks, or edit the model's source

Faced with this, developers reach for whatever the framework provides: PyTorch's module hooks, TensorFlow's session hooks, or, when those fall short, wrapping the optimizer's step function or editing the model's source code directly. Each task is implemented against these interfaces one at a time, extracting the tensors it needs and inserting its logic at the points it cares about.
Why it fell short

The hooks are fragmented across tasks, tensors, backends, and modes

This fragments badly. Different tasks need different injection points — pruning weights acts every training iteration, activation pruning and quantization act at each layer — and different tensors — weights, activations, or the hidden backward gradients. The frameworks expose these inconsistently, and worse, differently across their two execution modes: graph mode compiles the whole model first, while eager mode runs operators immediately. So each tool ends up bound to one backend and often one model, demanding low-level framework knowledge, and cannot be reused or composed.
This paper's idea

Borrow instrumentation from program analysis, at the operator level

Amanda observes that this is exactly the problem program-analysis tools like Pin already solved: they let you insert analysis code into a program without touching its source, reusable across programs. Amanda brings that to neural networks at the level of operators — the one abstraction every auto-differentiation framework shares, since all of them represent a model as a graph of operators. Write a tool once against a single before-and-after-operator interface, with all tensors reachable as context, and run it unchanged on any model, in either backend, kept cheap by caching.

Why this problem is worth solving

The heart of the mess is that a single conceptual task — “reach into the model here and read or change this” — has no single implementation, because the frameworks scatter it across many axes.

Consider the ways a task must specialize. It must choose an instrumentation point: weight pruning modifies static weight tensors once per training iteration, while activation pruning and quantization act inside each layer, at the operator, on runtime tensors. It must choose which computation state to touch: static post-training quantization needs only the weights, dynamic quantization needs weights and activations, and quantization-aware training additionally needs the gradient tensors from back-propagation. And it must contend with the execution mode: graph mode compiles the model into a computation graph up front, so the global structure is visible but any change means rewriting the graph, whereas eager mode executes each operator immediately, with no explicit graph but with complex control flow. The frameworks offer different, partial interfaces for each of these — a module hook here, a session hook there, source wrapping when neither fits — and PyTorch’s built-in quantization only covers basic convolution and linear layers, breaking on models that use functional calls. The result is that every research tool is a bespoke integration against a wide, inconsistent surface, tied to one framework and frequently one model, requiring its author to understand framework internals that have nothing to do with their actual idea. Multiply that across the hundreds of analysis and optimization techniques being invented, and the field is re-solving the same plumbing problem over and over, which is exactly the waste a unifying abstraction could end.

The reason a clean solution is possible, and why program analysis is the right lens, is that a good abstraction already exists — it just has not been named.

Program-analysis frameworks like Pin and NVBit succeeded by finding the right universal unit: the instruction. You write an analysis once against “before this instruction, after this instruction,” and it runs on any program, because every program is ultimately a stream of instructions. Amanda’s key observation is that neural networks have an equally universal unit that no one had claimed as an instrumentation point: the operator. Every modern auto-differentiation framework, however different its operator set, represents a model the same way — as a data-flow graph whose nodes are operators (multi-input, multi-output functions) and whose edges are tensors, with automatic differentiation defined per operator in reverse. That shared structure means the operator boundary is available everywhere, and it is the finest granularity these tasks actually need. Everything coarser — a module, a subgraph, a whole forward or backward pass, a training iteration — can be derived by placing operator-level points and combining them with context. And because tensors are just the edges of the graph, every computation state a task might want — a static weight, a runtime activation, an implicit backward gradient — is reachable through one interface. So the sprawling surface collapses to a single question, “what do you want to do before or after each operator, and which of its tensors do you need,” which is what makes a unified framework not just desirable but achievable.

What people did before, and where it broke

Use the framework’s module or session hooks. Register callbacks through PyTorch’s module hooks or TensorFlow’s session hooks. Where it broke. They cover only a fixed set of coarse points and states — PyTorch’s built-in quantization handles only basic convolution and linear layers, and models using functional calls break module-wrapping entirely.
Wrap the optimizer or edit the model’s source code. When the hooks are too limited, manually wrap the training step or rewrite the network’s source to inject the needed logic. Where it broke. The tool becomes tied to specific models and must be re-edited for each new one, so it does not scale as the number of supported networks grows.
Write separate implementations per execution mode. Handle graph mode and eager mode with their own code paths, since their interfaces differ. Where it broke. The same conceptual task needs two incompatible implementations, doubling the work and the framework knowledge required, and the two rarely stay in sync.
Rely on a purpose-built tool for one pattern. Use a dedicated library that supports one pruning or quantization pattern. Where it broke. It supports only its fixed pattern and often only one model API, so users must refactor their model to fit it, and it cannot compose with other tools.
The gap. Every approach forced developers to work against a wide, inconsistent surface of points, tensors, backends, and modes, producing tools bound to one setting that could not be reused or composed. What was missing was a single abstraction that unifies where you instrument and what you read, portable across models and backends — which is exactly what the flood of DNN analysis and optimization work needs.

How Amanda works

One operator-level instrumentation abstraction over every backend

↓ →
new
↓ →
new
↓ →
new
↓ →
new
↓ →
new

The idea

Import a mature idea from program analysis rather than inventing one. The field had been treating DNN analysis and optimization as a fresh plumbing problem for each task, when program analysis had long ago crystallized the same shape of problem into instrumentation: insert your logic around a program’s units without touching its source, and reuse it across programs. Amanda recognizes the neural-network version of that unit is the operator, and lifts the whole instrumentation discipline over. Borrowing a proven abstraction from an adjacent domain — inserting analysis at a natural execution boundary without editing the target — is exactly what the far-memory compiler work covered elsewhere on this site does, adopting binary-instrumentation-style guards to intercept memory accesses transparently rather than reinventing the mechanism from scratch.

Collapse a fragmented surface into one abstraction so a tool is written once. The pain was never any single hook but the sprawl of them across points, tensors, backends, and modes, forcing every tool into a bespoke integration. Amanda’s move is to prove that one operator-centric interface subsumes all of it — any point derivable from operator boundaries, any tensor reachable as a graph edge, any backend hidden behind a driver — so the developer writes against the task, not the framework, and the same tool ports to every model and both backends. Replacing a mass of ad-hoc, per-case, framework-specific effort with a single unified abstraction that a developer targets once is the same consolidation the recommendation-model compiler work covered elsewhere on this site achieves, turning what teams hand-optimized model by model into one automatic, model-independent path.

Make the general mechanism cheap by caching what does not change. A universal instrumentation layer that patched every operator on every run would be too slow to use, so its practicality hinges on paying the cost only where instrumentation actually happens. Amanda caches aggressively: operators no tool touches revert to their plain form, expensive analysis setup runs once and is reused, and the instrumented graph is reused until a tool modifies it. Determining what is stable and reusing it, recomputing only the part a change actually disturbs, is the same discipline the dynamic-neural-network compiler work covered elsewhere on this site applies, holding fixed everything it can settle ahead of time and re-doing only what the specific situation forces.

Following a pruning tool through Amanda

Writing one pruning tool and running it everywhere

Tracing a pruning tool shows how one Amanda tool replaces many ad-hoc ones.

Step 1 of 5

A researcher has a pruning method but many target models

The method decides which weights or activations to zero out during training. Traditionally the researcher would write it against one framework's hooks and hand-edit each model's source to reach the right tensors and splice into the training loop — code that grows with every new model supported.

Real community pruning projects run to hundreds or over a thousand lines for exactly this reason.

Why it works

The operator is the one point every framework shares, so a single interface fits every task. Each task wants a different injection point, from a whole iteration down to a single layer → but all auto-differentiation frameworks model a network as a graph of operators, and coarser points derive from operator boundaries → so one before-and-after-operator abstraction covers tracing, profiling, pruning, and quantization alike.
Every tensor is a graph edge, so one context reaches all computation state. Tasks need different tensors — static weights, runtime activations, hidden backward gradients → but every tensor is an edge connecting operators in the same graph → so all of them are reachable through a single context interface instead of task-specific extraction.
A backend-independent core over thin drivers hides the framework differences. Backends and execution modes expose inconsistent, mode-specific hooks → but the operator abstraction sits in a core above per-backend drivers that translate to each native mechanism → so a tool written once runs unchanged on PyTorch eager and TensorFlow graph.
Caching makes patching every operator affordable. Instrumenting a general layer that touches every operator on every run would be ruinously slow → but operators no tool instruments revert to plain execution and the instrumented graph is reused until it changes → so overhead stays under a few percent, with caching alone giving up to seventy-two times speedup.
Decoupling the tool from the model makes it portable and composable. Hand-written tools mix task logic with model and execution code, binding them to one setting → but an Amanda tool references only operators and tensors, and tools can layer on one another → so one tool ports to any model with up to six times less code and composes with others.

What it costs (trade-offs and scope)

Operator granularity is the floor, not kernel or binary level. Amanda instruments at operators and coarser subgraphs; finer instrumentation of individual GPU kernels or binary instructions is out of scope, though it composes cleanly with existing tools like CUPTI, NVBit, and Pin for that layer.
Overhead is low but uneven. It stays under one percent in eager mode and under about seven percent in graph mode for most cases, but forward-backward subgraph matching makes some models (like BERT) costlier in eager mode, many-operator models (like Inception) costlier in graph mode, and a small Transformer in TensorFlow at batch sizes of one or two can exceed ten percent memory overhead because the fixed management cost is relatively large.
Only two backends and modes are implemented. Drivers exist for PyTorch eager and TensorFlow graph; a new backend needs a new driver, and while the backend interface is clean, that is real work.
The drivers hide genuine complexity that must be gotten right. Patching every operator, snooping the registration process, mapping backward operators to their forward operators, and disabling differentiation on injected code so it does not corrupt the backward graph are subtle, backend-specific engineering.
Cross-backend portability needs an extra step. A tool is portable across models within a backend by default, but making it portable across different backends requires the additional context-mapping tool that transforms the backend’s context in advance.
It is infrastructure, not a new algorithm. Amanda unifies how analyses and optimizations are implemented, not what they do; its value depends on adoption and on the community contributing and composing tools over time, and it reproduces existing techniques rather than inventing new ones.

The evaluation

Three ways to instrument a neural network

Property	Framework hooks	Source modification
Where you can instrument
Tensors you can reach	Mostly module-level	Any, extracted manually
Portability	Tied to one backend
Execution modes	One mode, separate code	Per-mode effort
Code to build a tool	Grows with each model	Most, grows with models

Three ways to instrument a neural network, evaluated by reproducing community analysis and optimization tools as Amanda tools. Amanda implements tracing, profiling, interpretability, pruning, and quantization through one operator-level abstraction across PyTorch and TensorFlow, cutting five pruning projects from hundreds or over a thousand lines to a fraction (up to about six times fewer) at the same model accuracy, while keeping overhead under 5 percent in most cases and gaining up to 72 times from its caching — because a single operator abstraction with a unified tensor context replaces the fragmented, per-model, per-backend hooks that force tools to be rewritten by hand.Underlined values explain themselves when clicked.

Where this stops working

What to keep in mind

Amanda is a clean piece of infrastructure that names and solves a real, pervasive annoyance: every tool that studies or shrinks a neural network reimplements the same plumbing against fragmented, framework-specific hooks. Its insight is that program analysis already solved this shape of problem with instrumentation, and that the neural-network equivalent of the instruction is the operator — the one unit every auto-differentiation framework shares. By exposing before-and-after-operator points and every tensor as a unified context, over a backend-independent core with per-backend drivers, and caching aggressively, it lets a tool be written once and run unchanged on any model in PyTorch or TensorFlow, cutting community pruning projects up to six times in code at the same accuracy, with under 5 percent overhead in most cases. Read it with its scope. It works at operator granularity, deferring kernel and binary instrumentation to existing tools it composes with; overhead climbs for backward-subgraph-heavy models, many-operator models, and tiny inputs; only PyTorch eager and TensorFlow graph have drivers, and a new backend or cross-backend portability takes extra work; and it is an infrastructure contribution whose payoff depends on adoption, unifying how techniques are built rather than inventing them. The durable lesson is that when a community keeps rebuilding the same integration by hand, the highest-leverage move is often to find the universal unit hiding in plain sight and build one abstraction on it.

Background

Why maturing fields grow shared instrumentation layers

It helps to see Amanda as a sign of a field growing up. Early in any computing discipline, everyone who wants to observe or modify a running system does it by hand, poking at internals in whatever ad-hoc way works, because the system is changing too fast and the community is too small to justify shared tooling. As the field matures and the number of people doing similar things explodes, that hand-work becomes the bottleneck, and someone finds the right abstraction that lets a tool be written once and shared — the way binary instrumentation frameworks did for program analysis decades ago, turning a thousand bespoke profilers into reusable tools over a common interface. Deep learning has reached that point: the sheer volume of analysis and optimization research means the same reach-into-the-model plumbing is being rebuilt constantly, and the leverage now lies in shared infrastructure rather than one more clever technique. The deep-learning-on-GPUs retrospective covered elsewhere on this site makes a kindred point about the field’s maturation, arguing that sound progress increasingly comes from measuring and understanding what these systems actually do rather than assuming — and measurement at scale is precisely what a general instrumentation layer enables. The broad lesson is that the tools which most accelerate a maturing field are often not new capabilities but new abstractions that let the community stop re-solving the same problem and start building on each other’s work.

If you remember only three things

1. A vast amount of deep-learning research reaches into running models — tracing neurons, profiling compute, extracting interpretable substructures, pruning, quantizing — and all of it needs to monitor and manipulate a model's execution. But frameworks offer only fragmented, ad-hoc support: each task needs different injection points (a training iteration, a layer, an operator) and different tensors (weights, activations, gradients), and PyTorch and TensorFlow expose these through inconsistent hooks that also differ between graph mode (compile-then-run) and eager mode (run immediately). So every tool is hand-built against a wide surface, bound to one backend and often one model, and cannot be reused or composed.
2. Amanda borrows the instrumentation concept from program analysis (Pin, NVBit) and applies it at the operator level — the universal unit, since every auto-differentiation framework represents a model as a data-flow graph of operators with tensors as edges. A single user-level interface lets a tool register logic before or after any operator (forward and backward) and reach any tensor through a unified context, so coarser points and any computation state derive from that one abstraction. A backend-independent core over thin per-backend drivers (PyTorch eager, TensorFlow graph) hides framework differences by patching all operators, mapping backward operators to their forward ops, and preventing injected code from corrupting auto-differentiation, while aggressive caching (reverting un-instrumented operators to plain execution, reusing the instrumented graph) keeps the general mechanism cheap.
3. Reproducing five community pruning projects as Amanda tools cut their code from hundreds or over a thousand lines to a fraction — up to about six times fewer — at the same model accuracy, and Amanda also implements tracing, profiling, interpretability, and quantization through the one abstraction, composing with kernel-level tools like CUPTI. Overhead is under 1 percent in eager mode and under about 7 percent in graph mode for most cases, with caching giving up to 72 times speedup. Its limits: operator granularity is the floor (kernel/binary is deferred to other tools), overhead climbs for backward-subgraph-heavy or many-operator models and tiny inputs, only two backends and modes have drivers, cross-backend portability needs a context-mapping tool, and it is infrastructure whose value depends on community adoption rather than a new algorithm.

Check yourself

Check your understanding

1. Why is implementing DNN analysis and optimization tasks so fragmented in current frameworks?
2. What is the key abstraction Amanda uses to unify DNN instrumentation?
3. How does Amanda keep the overhead of instrumenting every operator low?
4. Which statement best reflects the scope and limits of Amanda?

Sources

The paper is “Amanda: Unified Instrumentation Framework for Deep Neural Networks” by Yue Guan, Yuxian Qiu, Jingwen Leng, Fan Yang, Shuo Yu, Yunxin Liu, Yu Feng, Yuhao Zhu, Lidong Zhou, Yun Liang, Chen Zhang, Chao Li, and Minyi Guo, published at ASPLOS 2024, DOI 10.1145/3617232.3624864. The authors are at Shanghai Jiao Tong University, the Shanghai Qi Zhi Institute, Microsoft Research, Tsinghua University, the Shanghai Artificial Intelligence Laboratory, the University of Rochester, and Peking University, and Amanda is open-source.

This explainer is written from the full 18-page text. Everything above comes from the paper’s own sections: the introduction and motivation establishing the pains of ad-hoc instrumentation points (weight pruning per iteration versus operator-level activation pruning), fragmented state representations (static, dynamic, and quantization-aware training needing weights, activations, and gradients), and execution-mode complexity (graph versus eager); the operator instrumentation abstraction drawing the analogy to binary instrumentation frameworks like Pin and NVBit and arguing the operator is the universal, finest-granularity point since all auto-differentiation frameworks share a data-flow-graph-of-operators design with tensors as edges; the Amanda interface with its user-level tool APIs (registration for analysis routines before and after forward and backward operators, instrumentation for on-the-fly modification, inspection for context, and control for scope and caches) and the analysis-versus-instrumentation-routine distinction; the two-layer system design separating the backend-independent Amanda-core from backend drivers for PyTorch eager and TensorFlow graph, including operator patching by snooping registration, backward-to-forward operator mapping, disabling auto-differentiation on injected code, consistent operator identifiers, and the eager action cache and graph switching-and-caching mechanisms; and the evaluation reproducing pruning projects (Tile-Wise, Dynamic Channel, Activation, Attention, and APEX Vector-Wise pruning) at the same accuracy with large code reductions (for example Tile-Wise from 1203 to 213 lines, Attention from 1105 to 179), demonstrating tracing, profiling, effective-path interpretability, and quantization, composing with CUPTI for kernel-level profiling, and reporting overhead under 1 percent in eager mode and under about 7 percent in graph mode, memory overhead under 5 percent in most cases, and up to 72.6 times (17.1 times average) speedup from caching.