© 2026 Nokia, Licensed under the BSD 3-Clause Clear License, SPDX-License-Identifier: BSD-3-Clause-Clear
A task-independent, layer-wise model for estimating the inference energy of neural networks.
WattLayer estimates the energy a neural network consumes during inference by modelling each layer type independently and summing the per-layer predictions. Because layers are shared across architectures and tasks, the same per-layer models compose to estimate the energy of unseen architectures — and even unseen tasks — directly from a model's description, without re-measuring on hardware. It ships three independent, YAML-driven pipelines — data collection, training, and estimation — and a first-class architecture-extraction API you can reuse on any PyTorch model.
This is the reference implementation of the SuRE'26 (IJCAI 2026) paper WattLayer: Get Layers Right to Estimate Inference Energy of Neural Networks (see Citation).
The widespread adoption of AI has raised concerns about its energy footprint, yet there is no standardized way to accurately estimate inference energy across tasks and architectures. WattLayer addresses this with a task-independent, layer-wise estimator:
- Validated at scale — evaluated on more than 100,000 layers from 295 neural network architectures, across 3 widely used tasks and 3 distinct hardware platforms.
- State-of-the-art accuracy — a median error of 19.6%, outperforming prior methods.
- Generalizes to new tasks — layer-wise decomposition transfers to new tasks without complete retraining, by reusing layers shared across architectures.
- Practical tooling — an open, reproducible methodology plus the tools and insights to help stakeholders design more energy-efficient AI systems.
WattLayer addresses the emerging need for methodologies and tools that enable users to predict the energy consumption of AI services proactively. It provides accurate predictions of neural network energy consumption at the layer level, enabling stakeholders to make informed decisions about energy efficiency during AI development and deployment. WattLayer learns the relationship between a layer's structural features (activations, MACs, parameters, input shape, …) and its measured energy, then composes those per-layer models to estimate the energy of unseen architectures from their description alone. Crucially, the layer-wise decomposition is task-independent: since architectures across tasks reuse the same kinds of layers, models trained on existing tasks transfer to new ones without complete retraining. Two whole-model baselines (Hidden-Joules and a MAC-only model) are kept for comparison.
WattLayer targets Python ≥ 3.10.
git clone <your-fork-url>
cd WattLayer
pip install -e ".[dev]" # runtime + dev tooling (ruff, mypy, pytest)The core runtime depends on torch, pandas, scikit-learn, ptflops,
transformers, timm, huggingface_hub, and codecarbon. Data collection
additionally needs a CUDA-capable GPU and the
codecarbon RAPL permissions; training and
estimation run on CPU.
WattLayer/
configs/ # one YAML per pipeline
data_collection.yaml
training.yaml
estimation.yaml
wattlayer/
architecture/ # PUBLIC API: extract_architecture, count_module_macs, create_model
utils/ # codecarbon parsing, shape helpers, logging
data_collection/ # DataCollector ABC, model sources (csv | huggingface), pipeline
training/ # dataset building, regression trainers, training pipeline
estimation/ # energy estimators + estimation pipeline
run_data_collection.py # thin CLI -> wattlayer.data_collection
run_train_models.py # thin CLI -> wattlayer.training
run_estimation.py # thin CLI -> wattlayer.estimation
examples/ # architecture-extraction usage examples
data/ # minimal example data + one example trained model
tests/ # mirrors wattlayer/ (no GPU, no network)
# 1. Collect energy measurements (requires a GPU)
python run_data_collection.py --config configs/data_collection.yaml
# 2. Train the layer-wise, Hidden-Joules, and MAC models
python run_train_models.py --config configs/training.yaml
# 3. Estimate energy for a set of models
python run_estimation.py --config configs/estimation.yamlA ready-to-load example model called WattLayer_Full_Training_Dataset_2026-06-01.pkl lives in models/, so you can try
estimation without retraining by pointing models_dir at it.
Runs forward passes for each model and layer while codecarbon records GPU
energy, writing a raw measurements CSV. Models come from either an explicit CSV
list or HuggingFace discovery (source: csv | huggingface). Requires a GPU.
Loads the raw measurements, attaches per-layer architecture features, and trains
one model per layer type (the layer-wise WattLayer model) plus the Hidden-Joules
and MAC baselines. Writes layer_models.pkl, hidden_joules_model.pkl,
mac_model.pkl, and a metrics.json summary. CPU-only.
Loads the trained models, extracts each target model's architecture (no training), estimates energy with all three models, and writes a results CSV. Optionally computes relative errors against a measured-energy CSV. CPU-only.
Each pipeline is fully described by a YAML file in configs/; there are no
hardcoded paths or magic constants in the code. See the inline comments in
configs/data_collection.yaml,
configs/training.yaml, and
configs/estimation.yaml for every option.
Both data collection and estimation share a models: block:
models:
source: csv # explicit list
csv_path: "data/models.csv"
# --- or ---
source: huggingface # discover most-downloaded models per task
tasks: ["image-classification"]
limit: 30
task_and_library: "data/models_task_and_library.csv"extract_architecture is the reusable, documented entry point for turning any
PyTorch model into a per-layer feature table. Point colleagues here.
import torch
from wattlayer.architecture import extract_architecture
model = torch.nn.Sequential(
torch.nn.Conv2d(3, 16, 3, padding=1),
torch.nn.ReLU(),
torch.nn.AdaptiveAvgPool2d((1, 1)),
torch.nn.Flatten(),
torch.nn.Linear(16, 10),
)
info = extract_architecture(model, example_input=torch.randn(1, 3, 32, 32), include_macs=True)
df = info.to_dataframe()
# Per-layer rows: model_name, layer_idx, module_type, kernel_size, in/out_channels,
# stride, padding, input_shape, activations, macs, parameters
# (row 0 is the "full_architecture" aggregate).You can also load by name (uses the network) via
extract_architecture_from_pretrained(...).
See examples/extract_architecture_example.py
for a fully offline walkthrough on a tiny model.
- Data collection: a raw
codecarbonCSV (one row per layer per repeat). - Training:
layer_models.pkl,hidden_joules_model.pkl,mac_model.pkl, andmetrics.json. - Estimation: a results CSV with
Energy_Layer,Energy_HJ,Energy_Macper model (plus relative-error columns when measured energy is supplied).
If you use WattLayer in academic work, please cite:
Adrien Sardi, Marie Line Alberi-Morel, Sara Alouf, Frédéric Giroire, and Joanna Moulierac. WattLayer: Get Layers Right to Estimate Inference Energy of Neural Networks. In SuRE'26: 1st Workshop on Sustainability and Resource-Efficiency of Artificial Intelligence, co-located with IJCAI 2026, August 17, 2026, Bremen, Germany.
@inproceedings{sardi2026wattlayer,
title = {WattLayer: Get Layers Right to Estimate Inference Energy of Neural Networks},
author = {Sardi, Adrien and Alberi-Morel, Marie Line and Alouf, Sara and Giroire, Fr\'ed\'eric and Moulierac, Joanna},
booktitle = {SuRE'26: 1st Workshop on Sustainability and Resource-Efficiency of Artificial Intelligence, co-located with IJCAI 2026},
address = {Bremen, Germany},
year = {2026},
month = aug,
}The data-collection machinery is adapted from
dl-energy-estimator
by Getzner et al. The estimation models, the layer-wise methodology, and this
packaged refactor are original to WattLayer.