Custom Architectures
This tutorial covers how to create custom neural network architectures in PyNAS, including defining custom blocks, creating specialized architectures, and integrating them into the evolutionary process.
Overview
PyNAS provides a flexible framework for defining custom architectures through:
Custom Blocks: Define new building blocks for your architectures
Architecture Templates: Create structured architecture patterns
Block Vocabulary: Manage available blocks for evolution
Constraints: Apply architectural constraints during evolution
Creating Custom Blocks
Custom Convolution Block
Here’s how to create a custom convolution block with specific properties:
import torch
import torch.nn as nn
from pynas.blocks.convolutions import ConvBlock
class DepthwiseSeparableConv(nn.Module):
"""
Custom depthwise separable convolution block.
Args:
in_channels (int): Number of input channels
out_channels (int): Number of output channels
kernel_size (int): Kernel size for convolution
stride (int): Stride for convolution
padding (int): Padding for convolution
"""
def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1):
super().__init__()
self.depthwise = nn.Conv2d(
in_channels, in_channels, kernel_size, stride, padding, groups=in_channels
)
self.pointwise = nn.Conv2d(in_channels, out_channels, 1)
self.bn1 = nn.BatchNorm2d(in_channels)
self.bn2 = nn.BatchNorm2d(out_channels)
self.relu = nn.ReLU(inplace=True)
def forward(self, x):
x = self.relu(self.bn1(self.depthwise(x)))
x = self.relu(self.bn2(self.pointwise(x)))
return x
Custom Attention Block
Create an attention mechanism for your architectures:
class SpatialAttention(nn.Module):
"""
Spatial attention mechanism for feature enhancement.
Args:
channels (int): Number of input channels
reduction (int): Channel reduction ratio
"""
def __init__(self, channels, reduction=16):
super().__init__()
self.avg_pool = nn.AdaptiveAvgPool2d(1)
self.max_pool = nn.AdaptiveMaxPool2d(1)
self.fc = nn.Sequential(
nn.Conv2d(channels, channels // reduction, 1, bias=False),
nn.ReLU(inplace=True),
nn.Conv2d(channels // reduction, channels, 1, bias=False)
)
self.sigmoid = nn.Sigmoid()
def forward(self, x):
avg_out = self.fc(self.avg_pool(x))
max_out = self.fc(self.max_pool(x))
attention = self.sigmoid(avg_out + max_out)
return x * attention
Integrating Custom Blocks
Register with Vocabulary
Add your custom blocks to the PyNAS vocabulary:
from pynas.core.vocabulary import Vocabulary
# Create vocabulary instance
vocab = Vocabulary()
# Register custom blocks
vocab.add_block('depthwise_separable', DepthwiseSeparableConv)
vocab.add_block('spatial_attention', SpatialAttention)
# Define block parameters
vocab.set_block_params('depthwise_separable', {
'kernel_size': [3, 5, 7],
'stride': [1, 2]
})
vocab.set_block_params('spatial_attention', {
'reduction': [8, 16, 32]
})
Custom Architecture Template
Create a template for generating architectures with your custom blocks:
from pynas.core.architecture_builder import ArchitectureBuilder
class CustomArchitectureBuilder(ArchitectureBuilder):
"""
Custom architecture builder with specific patterns.
"""
def __init__(self, vocab, input_channels=3, num_classes=10):
super().__init__(vocab)
self.input_channels = input_channels
self.num_classes = num_classes
def create_encoder_block(self, in_ch, out_ch, block_type='depthwise_separable'):
"""Create encoder block with optional attention."""
blocks = []
# Add main convolution
if block_type == 'depthwise_separable':
blocks.append(DepthwiseSeparableConv(in_ch, out_ch))
else:
blocks.append(nn.Conv2d(in_ch, out_ch, 3, padding=1))
blocks.append(nn.BatchNorm2d(out_ch))
blocks.append(nn.ReLU(inplace=True))
# Add attention if specified
blocks.append(SpatialAttention(out_ch))
return nn.Sequential(*blocks)
def build_architecture(self, genome):
"""Build architecture from genome representation."""
layers = []
current_channels = self.input_channels
for i, gene in enumerate(genome):
block_type = gene['type']
out_channels = gene['channels']
layer = self.create_encoder_block(
current_channels, out_channels, block_type
)
layers.append(layer)
current_channels = out_channels
# Add final classifier
layers.append(nn.AdaptiveAvgPool2d(1))
layers.append(nn.Flatten())
layers.append(nn.Linear(current_channels, self.num_classes))
return nn.Sequential(*layers)
Evolutionary Configuration
Configure evolution with custom architectures:
from pynas.core.population import Population
from pynas.core.individual import Individual
from pynas.opt.evo import GeneticAlgorithm
# Define custom genome structure
def create_custom_genome():
"""Create genome for custom architecture."""
genome = []
channels = [32, 64, 128, 256]
for ch in channels:
gene = {
'type': np.random.choice(['depthwise_separable', 'conv']),
'channels': ch,
'kernel_size': np.random.choice([3, 5, 7]),
'use_attention': np.random.choice([True, False])
}
genome.append(gene)
return genome
# Initialize population with custom genomes
population = Population(
population_size=20,
genome_factory=create_custom_genome,
architecture_builder=CustomArchitectureBuilder(vocab)
)
# Configure genetic algorithm
ga = GeneticAlgorithm(
population=population,
mutation_rate=0.1,
crossover_rate=0.8,
selection_method='tournament'
)
Advanced Customization
Multi-Scale Architecture
Create architectures that process multiple scales:
class MultiScaleBlock(nn.Module):
"""
Multi-scale processing block with parallel paths.
"""
def __init__(self, in_channels, out_channels):
super().__init__()
mid_channels = out_channels // 4
self.scale1 = nn.Conv2d(in_channels, mid_channels, 1)
self.scale2 = nn.Conv2d(in_channels, mid_channels, 3, padding=1)
self.scale3 = nn.Conv2d(in_channels, mid_channels, 5, padding=2)
self.scale4 = nn.Sequential(
nn.MaxPool2d(3, stride=1, padding=1),
nn.Conv2d(in_channels, mid_channels, 1)
)
self.bn = nn.BatchNorm2d(out_channels)
self.relu = nn.ReLU(inplace=True)
def forward(self, x):
s1 = self.scale1(x)
s2 = self.scale2(x)
s3 = self.scale3(x)
s4 = self.scale4(x)
out = torch.cat([s1, s2, s3, s4], dim=1)
return self.relu(self.bn(out))
Conditional Architecture
Create architectures that adapt based on conditions:
class ConditionalBlock(nn.Module):
"""
Block that selects operation based on learned gating.
"""
def __init__(self, in_channels, out_channels, num_ops=3):
super().__init__()
self.num_ops = num_ops
# Define multiple operations
self.ops = nn.ModuleList([
nn.Conv2d(in_channels, out_channels, 3, padding=1),
DepthwiseSeparableConv(in_channels, out_channels),
MultiScaleBlock(in_channels, out_channels)
])
# Gating mechanism
self.gating = nn.Sequential(
nn.AdaptiveAvgPool2d(1),
nn.Flatten(),
nn.Linear(in_channels, num_ops),
nn.Softmax(dim=1)
)
def forward(self, x):
gates = self.gating(x).unsqueeze(-1).unsqueeze(-1)
outputs = []
for i, op in enumerate(self.ops):
outputs.append(gates[:, i:i+1] * op(x))
return sum(outputs)
Architecture Constraints
Apply constraints during evolution:
class ArchitectureConstraints:
"""
Define constraints for architecture evolution.
"""
def __init__(self, max_params=1e6, max_flops=1e9):
self.max_params = max_params
self.max_flops = max_flops
def check_constraints(self, individual):
"""Check if individual satisfies constraints."""
model = individual.phenotype
# Check parameter count
param_count = sum(p.numel() for p in model.parameters())
if param_count > self.max_params:
return False
# Check FLOPs (simplified estimation)
flops = self.estimate_flops(model)
if flops > self.max_flops:
return False
return True
def estimate_flops(self, model):
"""Estimate FLOPs for the model."""
# Simplified FLOP estimation
total_flops = 0
for module in model.modules():
if isinstance(module, nn.Conv2d):
# Rough FLOP estimation for conv layers
kernel_flops = module.kernel_size[0] * module.kernel_size[1]
output_elements = 224 * 224 # Assume input size
total_flops += kernel_flops * output_elements * module.in_channels * module.out_channels
return total_flops
Complete Example
Here’s a complete example putting it all together:
import torch
import torch.nn as nn
import numpy as np
from pynas.core.population import Population
from pynas.opt.evo import GeneticAlgorithm
def run_custom_nas():
"""Run NAS with custom architectures."""
# Setup vocabulary with custom blocks
vocab = setup_custom_vocabulary()
# Create custom architecture builder
builder = CustomArchitectureBuilder(vocab, input_channels=3, num_classes=10)
# Initialize population
population = Population(
population_size=20,
genome_factory=create_custom_genome,
architecture_builder=builder
)
# Setup constraints
constraints = ArchitectureConstraints(max_params=1e6, max_flops=1e9)
# Configure genetic algorithm
ga = GeneticAlgorithm(
population=population,
mutation_rate=0.1,
crossover_rate=0.8,
constraints=constraints
)
# Run evolution
for generation in range(50):
# Evaluate population
ga.evaluate_population()
# Apply constraints
ga.filter_by_constraints()
# Evolve
ga.evolve()
# Log progress
best_individual = ga.get_best_individual()
print(f"Generation {generation}: Best fitness = {best_individual.fitness}")
return ga.get_best_individual()
if __name__ == "__main__":
best_architecture = run_custom_nas()
print(f"Best architecture found: {best_architecture}")
Next Steps
After creating custom architectures:
Validation: Test your custom blocks thoroughly
Integration: Ensure compatibility with PyNAS framework
Performance: Profile your custom blocks for efficiency
Documentation: Document your custom components
Sharing: Consider contributing useful blocks back to PyNAS
See also
Advanced Configuration for advanced configuration options
<no title> for more custom block examples
Blocks Module for API reference on blocks