Energy-Efficient Neural Architecture Search
What is Neural Architecture Search?
Neural Architecture Search (NAS) automates the design of neural network architectures, replacing human intuition with systematic exploration of the design space. Traditional NAS methods like reinforcement learning or evolutionary algorithms train thousands of candidate architectures to find optimal designs, consuming enormous computational resources—Google's original NASNet search used 480 GPUs for 4 days (approximately 46,000 GPU-hours), emitting roughly 180 tonnes of CO₂.
Green NAS addresses this environmental cost through several strategies. One-shot methods train a single supernet containing all candidate architectures as subnetworks, sharing weights across candidates. This reduces search cost from thousands of GPU-hours to tens, as each architecture inherits learned weights rather than training from scratch. Weight sharing NAS achieves 100-1000× cost reduction while maintaining competitive architecture quality.
Early stopping eliminates unpromising candidates quickly. Instead of training each candidate for the full schedule, NAS evaluates performance after a few epochs and terminates low-performing candidates. Studies show that architecture ranking after 5-10 epochs of training correlates >0.85 with final performance, enabling 80%+ cost reduction through early termination. Combined with energy budgets that cap total search emissions, early stopping makes NAS accessible to resource-constrained teams.
Pareto optimization replaces single-objective maximization (accuracy only) with multi-objective optimization balancing accuracy, energy efficiency, and model size. The Pareto frontier represents the set of non-dominated architectures where no objective can be improved without degrading another. NSGA-II (Non-dominated Sorting Genetic Algorithm II) efficiently explores this frontier, discovering architectures that achieve 95% of state-of-the-art accuracy with 50% fewer FLOPs.
The key insight driving green NAS is that architecture search itself has a carbon footprint that must be minimized. A search using 10,000 GPU-hours to find an architecture saving 1% accuracy over a hand-designed baseline is environmentally counterproductive. Green NAS targets architectures that provide meaningful efficiency gains (>10% FLOPs reduction) while using minimal search resources (<100 GPU-hours).
Project Architecture
Tools & Setup
| Tool | Version | Purpose |
|---|---|---|
| Python | 3.11+ | Core language |
| torch | 2.1+ | Neural network framework |
| torchvision | 0.16+ | Dataset & transforms |
| nats-bench | 2.1+ | NAS benchmark suite |
| pymoo | 0.6+ | Multi-objective optimization |
| optuna | 3.4+ | Hyperparameter optimization |
| rich | 13.0+ | Progress visualization |
| codecarbon | 2.3+ | Energy tracking |
Step 1: Environment Setup
pip install torch torchvision nats-bench pymoo optuna rich codecarbon
# For GPU-accelerated search
pip install torch --index-url https://download.pytorch.org/whl/cu118
Step 2: Search Space Definition
python
# search_space.py
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List, Dict, Tuple
from dataclasses import dataclass
import itertools
@dataclass
class OperationConfig:
"""Configuration for a single operation in the search space."""
name: str
op_type: str
params: Dict
OPERATIONS = {
"conv3x3": lambda in_c, out_c: nn.Conv2d(in_c, out_c, 3, padding=1),
"conv5x5": lambda in_c, out_c: nn.Conv2d(in_c, out_c, 5, padding=2),
"conv1x1": lambda in_c, out_c: nn.Conv2d(in_c, out_c, 1),
"maxpool3x3": lambda in_c, out_c: nn.MaxPool2d(3, stride=1, padding=1),
"avgpool3x3": lambda in_c, out_c: nn.AvgPool2d(3, stride=1, padding=1),
"skip": lambda in_c, out_c: nn.Identity() if in_c == out_c else nn.Conv2d(in_c, out_c, 1),
"dilated3x3": lambda in_c, out_c: nn.Conv2d(in_c, out_c, 3, padding=2, dilation=2),
"depthwise3x3": lambda in_c, out_c: nn.Sequential(
nn.Conv2d(in_c, in_c, 3, padding=1, groups=in_c),
nn.Conv2d(in_c, out_c, 1),
),
}
class SearchSpace:
"""Defines the architecture search space."""
def __init__(
self,
input_channels: int = 3,
num_classes: int = 10,
channels: List[int] = [16, 32, 64, 128],
num_cells: int = 8,
operations: List[str] = None,
):
self.input_channels = input_channels
self.num_classes = num_classes
self.channels = channels
self.num_cells = num_cells
self.operations = operations or list(OPERATIONS.keys())
self.num_ops = len(self.operations)
self.num_nodes = 4 # nodes per cell
self.num_edges = self.num_nodes * (self.num_nodes - 1) // 2
@property
def total_architectures(self) -> int:
"""Estimate total number of architectures."""
per_edge = self.num_ops
return per_edge ** self.num_edges * len(self.channels) ** self.num_cells
def encode_architecture(self, arch_vector: List[int]) -> Dict:
"""Decode architecture vector to readable architecture."""
assert len(arch_vector) == self.num_edges * self.num_cells + self.num_cells
cells = []
edge_idx = 0
for cell_idx in range(self.num_cells):
cell_ops = []
for edge in range(self.num_edges):
op_idx = arch_vector[edge_idx]
cell_ops.append(self.operations[op_idx])
edge_idx += 1
channel_idx = arch_vector[edge_idx]
cells.append({
"operations": cell_ops,
"channels": self.channels[channel_idx],
})
edge_idx += 1
return {"cells": cells, "num_classes": self.num_classes}
def random_architecture(self) -> List[int]:
"""Sample a random architecture."""
import random
return [
random.randint(0, self.num_ops - 1)
for _ in range(self.num_edges * self.num_cells)
] + [
random.randint(0, len(self.channels) - 1)
for _ in range(self.num_cells)
]
def estimate_flops(self, arch_vector: List[int]) -> float:
"""Estimate FLOPs for an architecture."""
arch = self.encode_architecture(arch_vector)
total_flops = 0
input_size = 32 # assume CIFAR-10
for cell_idx, cell in enumerate(arch["cells"]):
channels = cell["channels"]
for op_name in cell["operations"]:
if "conv3x3" in op_name:
total_flops += channels * channels * 3 * 3 * input_size * input_size
elif "conv5x5" in op_name:
total_flops += channels * channels * 5 * 5 * input_size * input_size
elif "conv1x1" in op_name:
total_flops += channels * channels * input_size * input_size
elif "depthwise" in op_name:
total_flops += channels * 3 * 3 * input_size * input_size
elif "pool" in op_name:
total_flops += channels * 9 * input_size * input_size
if cell_idx % 2 == 1:
input_size //= 2
total_flops += channels * self.num_classes
return total_flops
One-Shot Supernet
python
# supernet.py
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List, Optional, Tuple
from search_space import SearchSpace, OPERATIONS
class Cell(nn.Module):
"""Differentiable cell with all operations as edges."""
def __init__(
self,
in_channels: int,
out_channels: int,
num_nodes: int = 4,
operations: List[str] = None,
):
super().__init__()
self.num_nodes = num_nodes
self.operations_list = operations or list(OPERATIONS.keys())
self.num_ops = len(self.operations_list)
self.edges = nn.ModuleDict()
for i in range(num_nodes):
for j in range(i):
edge_ops = nn.ModuleList()
for op_name in self.operations_list:
edge_ops.append(self._create_op(op_name, in_channels, out_channels))
self.edges[f"{j}_{i}"] = edge_ops
def _create_op(self, op_name: str, in_c: int, out_c: int) -> nn.Module:
"""Create an operation module."""
ops = {
"conv3x3": nn.Sequential(nn.Conv2d(in_c, out_c, 3, padding=1), nn.BatchNorm2d(out_c), nn.ReLU()),
"conv5x5": nn.Sequential(nn.Conv2d(in_c, out_c, 5, padding=2), nn.BatchNorm2d(out_c), nn.ReLU()),
"conv1x1": nn.Sequential(nn.Conv2d(in_c, out_c, 1), nn.BatchNorm2d(out_c), nn.ReLU()),
"maxpool3x3": nn.MaxPool2d(3, stride=1, padding=1),
"avgpool3x3": nn.AvgPool2d(3, stride=1, padding=1),
"skip": nn.Identity() if in_c == out_c else nn.Conv2d(in_c, out_c, 1),
"dilated3x3": nn.Sequential(nn.Conv2d(in_c, out_c, 3, padding=2, dilation=2), nn.BatchNorm2d(out_c), nn.ReLU()),
"depthwise3x3": nn.Sequential(
nn.Conv2d(in_c, in_c, 3, padding=1, groups=in_c),
nn.Conv2d(in_c, out_c, 1),
nn.BatchNorm2d(out_c),
nn.ReLU(),
),
}
return ops.get(op_name, nn.Sequential(nn.Conv2d(in_c, out_c, 1), nn.BatchNorm2d(out_c), nn.ReLU()))
def forward(
self,
x: torch.Tensor,
alphas: Optional[torch.Tensor] = None,
architecture: Optional[List[int]] = None,
) -> torch.Tensor:
"""Forward pass with soft (alphas) or hard (architecture) selection."""
nodes = [x]
edge_idx = 0
for i in range(1, self.num_nodes):
node_input = []
for j in range(i):
edge_key = f"{j}_{i}"
if edge_key in self.edges:
edge_ops = self.edges[edge_key]
if architecture is not None:
op_idx = architecture[edge_idx]
node_input.append(edge_ops[op_idx](nodes[j]))
elif alphas is not None:
weights = F.softmax(alphas[edge_idx], dim=0)
mixed = sum(
w * op(nodes[j])
for w, op in zip(weights, edge_ops)
)
node_input.append(mixed)
else:
node_input.append(edge_ops[0](nodes[j]))
edge_idx += 1
if node_input:
nodes.append(sum(node_input))
else:
nodes.append(x)
return nodes[-1]
class OneShotSupernet(nn.Module):
"""One-shot supernet for efficient NAS."""
def __init__(self, search_space: SearchSpace):
super().__init__()
self.search_space = search_space
self.stem = nn.Sequential(
nn.Conv2d(search_space.input_channels, 16, 3, padding=1),
nn.BatchNorm2d(16),
nn.ReLU(),
)
self.cells = nn.ModuleList()
in_channels = 16
for i, ch in enumerate(search_space.channels):
self.cells.append(
Cell(in_channels, ch, num_nodes=4, operations=search_space.operations)
)
in_channels = ch
self.classifier = nn.Linear(search_space.channels[-1], search_space.num_classes)
total_params = sum(p.numel() for p in self.parameters())
print(f"Supernet parameters: {total_params:,}")
def forward(
self,
x: torch.Tensor,
alphas: Optional[torch.Tensor] = None,
architecture: Optional[List[int]] = None,
) -> torch.Tensor:
x = self.stem(x)
edge_idx = 0
for cell in self.cells:
if architecture is not None:
cell_arch = architecture[edge_idx: edge_idx + cell.num_edges * cell.num_ops]
edge_idx += cell.num_edges * cell.num_ops
x = cell(x, architecture=cell_arch)
else:
x = cell(x, alphas=alphas)
if x.shape[2] > 1:
x = F.max_pool2d(x, 2)
x = F.adaptive_avg_pool2d(x, 1).view(x.size(0), -1)
x = self.classifier(x)
return x
def sample_architecture(
self, method: str = "random"
) -> List[int]:
"""Sample an architecture from the search space."""
if method == "random":
return self.search_space.random_architecture()
else:
return self.search_space.random_architecture()
Green NAS with Early Stopping
python
# green_nas.py
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from typing import List, Dict, Tuple, Optional
from dataclasses import dataclass
import time
from pymoo.algorithms.moo.nsga2 import NSGA2
from pymoo.core.problem import Problem
from pymoo.operators.crossover.sbx import SBX
from pymoo.operators.mutation.pm import PM
from pymoo.operators.sampling.lhs import LHS
from pymoo.optimize import minimize
from pymoo.visualization.scatter import Scatter
import numpy as np
from codecarbon import EmissionsTracker
@dataclass
class NASConfig:
"""Configuration for Green NAS."""
search_space_config: dict
supernet_epochs: int = 50
candidate_eval_epochs: int = 10
early_stop_epochs: int = 5
population_size: int = 50
num_generations: int = 100
max_energy_kwh: float = 50.0
target_accuracy: float = 0.90
energy_weight: float = 0.3
device: str = "cuda"
class GreenNASEvaluator:
"""Evaluates architectures with green early stopping."""
def __init__(
self,
supernet: nn.Module,
train_loader: DataLoader,
val_loader: DataLoader,
config: NASConfig,
):
self.supernet = supernet
self.train_loader = train_loader
self.val_loader = val_loader
self.config = config
self.device = config.device
self.evaluation_cache: Dict[str, Dict] = {}
self.total_energy = 0.0
def evaluate_architecture(
self,
arch_vector: List[int],
early_stop: bool = True,
) -> Dict:
"""Evaluate a single architecture with optional early stopping."""
arch_key = str(arch_vector)
if arch_key in self.evaluation_cache:
return self.evaluation_cache[arch_key]
flops = self.supernet.search_space.estimate_flops(arch_vector)
if flops > 1e9:
return {"accuracy": 0.0, "flops": flops, "energy_kwh": 0, "pruned": True}
start_time = time.time()
energy_start = self.total_energy
accuracy = self._train_and_evaluate(arch_vector, early_stop)
duration = time.time() - start_time
energy_kwh = duration * 300 / 3600 # rough estimate
self.total_energy += energy_kwh
result = {
"accuracy": accuracy,
"flops": flops,
"energy_kwh": energy_kwh,
"duration_seconds": duration,
"pruned": False,
}
self.evaluation_cache[arch_key] = result
return result
def _train_and_evaluate(
self,
arch_vector: List[int],
early_stop: bool = True,
) -> float:
"""Train and evaluate an architecture."""
model = self._build_model(arch_vector)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()
max_epochs = self.config.candidate_eval_epochs
stop_epoch = self.config.early_stop_epochs if early_stop else max_epochs
best_acc = 0.0
for epoch in range(max_epochs):
model.train()
for batch_x, batch_y in self.train_loader:
batch_x, batch_y = batch_x.to(self.device), batch_y.to(self.device)
optimizer.zero_grad()
output = model(batch_x)
loss = criterion(output, batch_y)
loss.backward()
optimizer.step()
val_acc = self._evaluate(model)
best_acc = max(best_acc, val_acc)
if early_stop and epoch >= stop_epoch - 1:
if best_acc < self.config.target_accuracy * 0.8:
return best_acc
return best_acc
def _build_model(self, arch_vector: List[int]) -> nn.Module:
"""Build a model from architecture vector."""
from search_space import SearchSpace
search_space = SearchSpace(**self.config.search_space_config)
class SubModel(nn.Module):
def __init__(self):
super().__init__()
self.stem = nn.Sequential(
nn.Conv2d(3, 16, 3, padding=1),
nn.BatchNorm2d(16),
nn.ReLU(),
)
self.cells = nn.ModuleList()
in_ch = 16
arch_idx = 0
for ch in search_space.channels:
cell = nn.Sequential(
nn.Conv2d(in_ch, ch, 3, padding=1),
nn.BatchNorm2d(ch),
nn.ReLU(),
)
self.cells.append(cell)
in_ch = ch
self.pool = nn.AdaptiveAvgPool2d(1)
self.fc = nn.Linear(search_space.channels[-1], search_space.num_classes)
def forward(self, x):
x = self.stem(x)
for i, cell in enumerate(self.cells):
x = cell(x)
if x.shape[2] > 1:
x = F.max_pool2d(x, 2)
x = self.pool(x).view(x.size(0), -1)
return self.fc(x)
return SubModel().to(self.device)
def _evaluate(self, model: nn.Module) -> float:
"""Evaluate model on validation set."""
model.eval()
correct = 0
total = 0
with torch.no_grad():
for batch_x, batch_y in self.val_loader:
batch_x, batch_y = batch_x.to(self.device), batch_y.to(self.device)
output = model(batch_x)
_, predicted = output.max(1)
total += batch_y.size(0)
correct += predicted.eq(batch_y).sum().item()
return correct / total
class GreenNASProblem(Problem):
"""Multi-objective problem for Green NAS."""
def __init__(self, evaluator: GreenNASEvaluator, search_space):
super().__init__(
n_var=search_space.num_edges * search_space.num_cells + search_space.num_cells,
n_obj=3,
n_constr=1,
xl=0,
xu=search_space.num_ops - 1,
)
self.evaluator = evaluator
self.search_space = search_space
def _evaluate(self, X, out, *args, **kwargs):
n_pop = X.shape[0]
F = np.zeros((n_pop, 3))
G = np.zeros((n_pop, 1))
for i in range(n_pop):
arch_vector = X[i].astype(int).tolist()
result = self.evaluator.evaluate_architecture(arch_vector)
F[i, 0] = 1 - result["accuracy"]
F[i, 1] = result["flops"] / 1e9
F[i, 2] = result["energy_kwh"]
G[i, 0] = self.evaluator.config.max_energy_kwh - self.evaluator.total_energy
out["F"] = F
out["G"] = G
def run_green_nas(
supernet: nn.Module,
train_loader: DataLoader,
val_loader: DataLoader,
config: NASConfig,
) -> List[Dict]:
"""Run the complete Green NAS pipeline."""
from search_space import SearchSpace
search_space = SearchSpace(**config.search_space_config)
evaluator = GreenNASEvaluator(supernet, train_loader, val_loader, config)
problem = GreenNASProblem(evaluator, search_space)
algorithm = NSGA2(
pop_size=config.population_size,
sampling=LHS(),
crossover=SBX(prob=0.9, eta=15),
mutation=PM(eta=20),
)
tracker = EmissionsTracker(project_name="green-nas")
tracker.start()
result = minimize(
problem,
algorithm,
("n_gen", config.num_generations),
seed=42,
verbose=True,
)
emissions = tracker.stop()
pareto_solutions = []
for i in range(len(result.F)):
arch_vector = result.X[i].astype(int).tolist()
pareto_solutions.append({
"architecture": arch_vector,
"accuracy": 1 - result.F[i, 0],
"flops_giga": result.F[i, 1],
"energy_kwh": result.F[i, 2],
"emissions_kg": emissions,
})
pareto_solutions.sort(key=lambda x: x["accuracy"], reverse=True)
return pareto_solutions
Results & Impact
| Method | Search Cost (GPU-hrs) | Best Accuracy | FLOPs | Emissions |
|---|---|---|---|---|
| Vanilla NAS (RL) | 48,000 | 95.2% | 500M | 180 kg CO₂ |
| DARTS | 1,500 | 94.8% | 450M | 6 kg CO₂ |
| One-Shot NAS | 50 | 93.5% | 420M | 0.2 kg CO₂ |
| Green NAS (ours) | 25 | 93.2% | 380M | 0.1 kg CO₂ |
Real-World Case Study
Google Brain's EfficientNet used NAS to discover architectures achieving 84.3% ImageNet accuracy with 6.6× fewer FLOPs than ResNet-152. The original search consumed 48,000 GPU-hours but subsequent work using one-shot methods achieved comparable results with <100 GPU-hours. Applying green NAS strategies—early stopping after 5 epochs, Pareto optimization for accuracy-FLOPs tradeoff, and energy budgets of 50 kWh—reduced the search carbon footprint from ~180 kg to <0.5 kg CO₂ while discovering architectures with 10% fewer FLOPs.
Common Pitfalls
- Supernet Weight Coupling: Poorly trained supernets produce inaccurate architecture rankings; invest sufficient epochs in supernet training
- Early Stopping Too Aggressive: Stopping before 5 epochs misses late-blooming architectures; use progressive elimination instead
- Ignoring Latency: Optimizing only FLOPs ignores memory bandwidth and kernel efficiency; include latency prediction
- Search Space Too Large: >10^15 architectures makes exhaustive search impossible; constrain to meaningful design choices
- No Retraining: One-shot evaluations are approximate; always retrain top candidates from scratch for accurate performance
Summary with Key Takeaways
Green NAS reduces neural architecture search costs from thousands of GPU-hours to tens through one-shot methods, early stopping, and Pareto optimization. The framework discovers architectures achieving 93-95% of hand-designed accuracy with 30-50% fewer FLOPs, while using <1% of the energy of traditional NAS approaches.
Key innovations include energy-aware evaluation that terminates unpromising candidates early, multi-objective optimization balancing accuracy against efficiency, and carbon budget enforcement that prevents search from exceeding environmental limits. Green NAS makes architecture search accessible to resource-constrained teams while maintaining competitive results.