Production Code Examples

Complete implementation of ResNet-18 inference on photonic hardware with INT4 quantization.

"""
Photonic ResNet-18 Image Classifier
-----------------------------------
Maps ResNet-18 to a photonic MVM accelerator with:
- 4-bit weight quantization
- 32x32 MVM tiles
- Noise-aware inference
"""
import numpy as np
from photonic import PhotonicMVM, NoiseModel, TileMapper
from torchvision.models import resnet18
import torch

class PhotonicResNet18:
    def __init__(self, tile_size=32, precision=4, optical_power=1e-3):
        self.tile_size = tile_size
        self.precision = precision

        # Load pre-trained model
        self.model = resnet18(pretrained=True)
        self.model.eval()

        # Map to photonic tiles
        self.mapper = TileMapper(tile_size, precision)
        self.photonic_layers = self._map_layers()

        # Set up noise model
        self.noise = NoiseModel(
            optical_power=optical_power,
            bandwidth=20e9,
            temperature=300,
            responsivity=0.9
        )

    def _map_layers(self):
        """Map all Conv2D and Linear layers to photonic tiles"""
        layers = {}
        for name, module in self.model.named_modules():
            if isinstance(module, torch.nn.Conv2d):
                W = module.weight.detach().numpy()
                # Reshape for im2col-based convolution
                W_flat = W.reshape(W.shape[0], -1)
                layers[name] = self.mapper.map_matrix(W_flat)
            elif isinstance(module, torch.nn.Linear):
                W = module.weight.detach().numpy()
                layers[name] = self.mapper.map_matrix(W)
        return layers

    def forward(self, image):
        """Run inference through photonic network"""
        # Preprocessing (electronic)
        x = self._preprocess(image)

        # Conv layers with im2col
        for layer_name, tiles in self.photonic_layers.items():
            if 'conv' in layer_name:
                x = self._photonic_conv(x, tiles)
                x = np.maximum(0, x)  # ReLU (electronic)
            elif 'fc' in layer_name:
                x = self._photonic_fc(x.flatten(), tiles)

        return x

    def _photonic_conv(self, x, tiles):
        """Photonic convolution using im2col"""
        # Apply im2col transformation
        x_col = im2col(x, kernel_size=3, stride=1, padding=1)

        # Run through tiles
        output = np.zeros((tiles['output_channels'], x_col.shape[1]))
        for i, tile_row in enumerate(tiles['mvm_tiles']):
            for j, tile in enumerate(tile_row):
                tile.set_noise_model(self.noise)
                start_row = i * self.tile_size
                end_row = min((i + 1) * self.tile_size, output.shape[0])
                start_col = j * self.tile_size

                x_slice = x_col[start_col:start_col + self.tile_size, :]
                output[start_row:end_row, :] += tile.forward(x_slice)

        return col2im(output, x.shape)

    def _photonic_fc(self, x, tiles):
        """Photonic fully-connected layer"""
        output = np.zeros(tiles['output_dim'])
        for i, tile_row in enumerate(tiles['mvm_tiles']):
            tile_output = np.zeros(self.tile_size)
            for j, tile in enumerate(tile_row):
                start = j * self.tile_size
                end = min((j + 1) * self.tile_size, len(x))
                if start < len(x):
                    tile_output += tile.forward(x[start:end])

            start_out = i * self.tile_size
            end_out = min((i + 1) * self.tile_size, len(output))
            output[start_out:end_out] = tile_output[:end_out - start_out]

        return output

    def benchmark(self):
        """Return performance metrics"""
        total_tiles = sum(
            len(t['mvm_tiles']) * len(t['mvm_tiles'][0])
            for t in self.photonic_layers.values()
        )
        total_macs = sum(
            t['output_dim'] * t['input_dim']
            for t in self.photonic_layers.values()
        )

        clock_rate = 10e9
        throughput = total_tiles * self.tile_size**2 * 2 * clock_rate / 1e12

        return {
            'total_tiles': total_tiles,
            'total_macs': total_macs,
            'throughput_tops': throughput,
            'latency_us': total_macs / (throughput * 1e12) * 1e6,
            'energy_per_inference_uj': total_macs * 15e-15 * 1e6
        }

# Usage
if __name__ == "__main__":
    accelerator = PhotonicResNet18(tile_size=32, precision=4)

    # Load test image
    from PIL import Image
    img = Image.open('test_image.jpg').resize((224, 224))
    img_array = np.array(img) / 255.0

    # Run inference
    output = accelerator.forward(img_array)
    prediction = np.argmax(output)
    print(f"Predicted class: {prediction}")

    # Print benchmarks
    metrics = accelerator.benchmark()
    print(f"Throughput: {metrics['throughput_tops']:.1f} TOPS")
    print(f"Latency: {metrics['latency_us']:.2f} μs")
    print(f"Energy: {metrics['energy_per_inference_uj']:.1f} μJ")

High-throughput inference engine for real-time video processing.

"""
Real-Time Photonic Inference Engine
-----------------------------------
Pipelined inference for video streams at 1000+ FPS
"""
import numpy as np
from collections import deque
from threading import Thread
import queue

class PhotonicInferenceEngine:
    def __init__(self, model_path, batch_size=8, num_pipelines=4):
        self.batch_size = batch_size
        self.num_pipelines = num_pipelines

        # Input/output queues
        self.input_queue = queue.Queue(maxsize=100)
        self.output_queue = queue.Queue(maxsize=100)

        # Create parallel pipelines
        self.pipelines = []
        for i in range(num_pipelines):
            mvm = PhotonicMVMPipeline(
                tile_size=32,
                precision=4,
                num_tiles=64
            )
            mvm.load_model(model_path)
            self.pipelines.append(mvm)

        self.running = False
        self.workers = []

    def start(self):
        """Start inference workers"""
        self.running = True
        for i, pipeline in enumerate(self.pipelines):
            worker = Thread(target=self._inference_worker, args=(i, pipeline))
            worker.daemon = True
            worker.start()
            self.workers.append(worker)

    def stop(self):
        """Stop all workers"""
        self.running = False
        for worker in self.workers:
            worker.join(timeout=1.0)

    def _inference_worker(self, worker_id, pipeline):
        """Worker thread for batch inference"""
        batch = []
        batch_ids = []

        while self.running:
            try:
                item = self.input_queue.get(timeout=0.01)
                batch.append(item['data'])
                batch_ids.append(item['id'])

                if len(batch) >= self.batch_size:
                    # Process batch
                    batch_array = np.stack(batch)
                    results = pipeline.forward_batch(batch_array)

                    # Send results
                    for i, result in enumerate(results):
                        self.output_queue.put({
                            'id': batch_ids[i],
                            'result': result,
                            'worker': worker_id
                        })

                    batch = []
                    batch_ids = []

            except queue.Empty:
                # Process partial batch if waiting too long
                if batch:
                    batch_array = np.stack(batch)
                    results = pipeline.forward_batch(batch_array)
                    for i, result in enumerate(results):
                        self.output_queue.put({
                            'id': batch_ids[i],
                            'result': result,
                            'worker': worker_id
                        })
                    batch = []
                    batch_ids = []

    def infer_async(self, frame_id, data):
        """Submit frame for async inference"""
        self.input_queue.put({'id': frame_id, 'data': data})

    def get_result(self, timeout=None):
        """Get inference result"""
        return self.output_queue.get(timeout=timeout)

    def get_throughput(self):
        """Estimate throughput based on pipeline config"""
        ops_per_tile = 32 * 32 * 2  # MACs per tile per cycle
        clock_rate = 10e9
        ops_per_pipeline = 64 * ops_per_tile * clock_rate
        return self.num_pipelines * ops_per_pipeline / 1e12  # TOPS


# Video processing example
if __name__ == "__main__":
    import cv2

    engine = PhotonicInferenceEngine(
        model_path='yolov5s_photonic.bin',
        batch_size=4,
        num_pipelines=4
    )
    engine.start()

    cap = cv2.VideoCapture('input_video.mp4')
    frame_id = 0
    results = {}

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        # Preprocess and submit
        frame_resized = cv2.resize(frame, (640, 640)) / 255.0
        engine.infer_async(frame_id, frame_resized)

        # Collect results (non-blocking)
        try:
            result = engine.get_result(timeout=0.001)
            results[result['id']] = result['result']
        except queue.Empty:
            pass

        frame_id += 1

    engine.stop()
    print(f"Processed {frame_id} frames")
    print(f"Throughput: {engine.get_throughput():.1f} TOPS")

Train neural networks with photonic hardware constraints in the loop.

"""
Quantization-Aware Training for Photonic Accelerators
------------------------------------------------------
Train models that account for:
- Weight quantization
- Activation quantization
- Photonic noise
"""
import torch
import torch.nn as nn
import torch.nn.functional as F

class PhotonicQuantizer(torch.autograd.Function):
    """Straight-through estimator for quantization"""
    @staticmethod
    def forward(ctx, x, bits, symmetric=True):
        if symmetric:
            qmin = -2**(bits-1)
            qmax = 2**(bits-1) - 1
            scale = x.abs().max() / qmax
        else:
            qmin = 0
            qmax = 2**bits - 1
            scale = (x.max() - x.min()) / qmax

        x_q = torch.clamp(torch.round(x / scale), qmin, qmax)
        return x_q * scale

    @staticmethod
    def backward(ctx, grad_output):
        return grad_output, None, None


class PhotonicLinear(nn.Module):
    """Linear layer with photonic constraints"""
    def __init__(self, in_features, out_features, weight_bits=4, act_bits=8):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(out_features, in_features) * 0.1)
        self.bias = nn.Parameter(torch.zeros(out_features))
        self.weight_bits = weight_bits
        self.act_bits = act_bits

        # Photonic noise parameters
        self.noise_std = 0.01

    def forward(self, x):
        # Quantize weights
        w_q = PhotonicQuantizer.apply(self.weight, self.weight_bits)

        # Quantize activations
        x_q = PhotonicQuantizer.apply(x, self.act_bits, False)

        # Compute with quantized values
        out = F.linear(x_q, w_q, self.bias)

        # Add photonic noise during training
        if self.training:
            noise = torch.randn_like(out) * self.noise_std * out.abs()
            out = out + noise

        return out


class PhotonicConv2d(nn.Module):
    """Conv2D layer with photonic constraints"""
    def __init__(self, in_channels, out_channels, kernel_size,
                 stride=1, padding=0, weight_bits=4, act_bits=8):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(
            out_channels, in_channels, kernel_size, kernel_size) * 0.1)
        self.bias = nn.Parameter(torch.zeros(out_channels))
        self.stride = stride
        self.padding = padding
        self.weight_bits = weight_bits
        self.act_bits = act_bits
        self.noise_std = 0.01

    def forward(self, x):
        w_q = PhotonicQuantizer.apply(self.weight, self.weight_bits)
        x_q = PhotonicQuantizer.apply(x, self.act_bits, False)

        out = F.conv2d(x_q, w_q, self.bias, self.stride, self.padding)

        if self.training:
            noise = torch.randn_like(out) * self.noise_std * out.abs()
            out = out + noise

        return out


class PhotonicMNISTNet(nn.Module):
    """Example network for MNIST with photonic constraints"""
    def __init__(self, weight_bits=4, act_bits=8):
        super().__init__()
        self.conv1 = PhotonicConv2d(1, 32, 3, padding=1,
                                     weight_bits=weight_bits, act_bits=act_bits)
        self.conv2 = PhotonicConv2d(32, 64, 3, padding=1,
                                     weight_bits=weight_bits, act_bits=act_bits)
        self.fc1 = PhotonicLinear(64 * 7 * 7, 128,
                                   weight_bits=weight_bits, act_bits=act_bits)
        self.fc2 = PhotonicLinear(128, 10,
                                   weight_bits=weight_bits, act_bits=act_bits)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2)
        x = x.flatten(1)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x


# Training loop
def train_photonic_model(model, train_loader, epochs=10, lr=0.001):
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()

    for epoch in range(epochs):
        model.train()
        total_loss = 0

        for batch_idx, (data, target) in enumerate(train_loader):
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

        print(f"Epoch {epoch+1}: Loss = {total_loss/len(train_loader):.4f}")

    return model


if __name__ == "__main__":
    from torchvision import datasets, transforms
    from torch.utils.data import DataLoader

    # Load MNIST
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

    # Train with different precisions
    for bits in [8, 4, 2]:
        print(f"\n--- Training with {bits}-bit weights ---")
        model = PhotonicMNISTNet(weight_bits=bits, act_bits=8)
        model = train_photonic_model(model, train_loader, epochs=5)

        # Export for photonic deployment
        torch.save(model.state_dict(), f'mnist_photonic_{bits}bit.pth')

Active thermal stabilization for maintaining weight accuracy.

"""
Thermal Feedback Control for Photonic Weights
----------------------------------------------
PID controller for ring resonator stabilization
"""
import numpy as np
from collections import deque

class ThermalController:
    def __init__(self, num_rings, target_wavelength=1550e-9):
        self.num_rings = num_rings
        self.target_wl = target_wavelength

        # PID gains (tuned for typical thermal response)
        self.Kp = 5.0
        self.Ki = 0.5
        self.Kd = 0.1

        # State variables
        self.integral = np.zeros(num_rings)
        self.prev_error = np.zeros(num_rings)
        self.heater_power = np.zeros(num_rings)

        # Limits
        self.max_heater_power = 50.0  # mW
        self.min_heater_power = 0.0

        # History for monitoring
        self.error_history = deque(maxlen=1000)
        self.power_history = deque(maxlen=1000)

    def update(self, current_wavelengths, dt=0.001):
        """
        Update heater powers based on wavelength error

        Args:
            current_wavelengths: Array of current resonance wavelengths
            dt: Time step in seconds

        Returns:
            heater_powers: Array of heater power commands (mW)
        """
        # Calculate error
        error = self.target_wl - current_wavelengths

        # PID terms
        P = self.Kp * error
        self.integral += error * dt
        self.integral = np.clip(self.integral, -10, 10)  # Anti-windup
        I = self.Ki * self.integral
        D = self.Kd * (error - self.prev_error) / dt

        # Update heater power
        delta_power = P + I + D
        self.heater_power += delta_power

        # Apply limits
        self.heater_power = np.clip(
            self.heater_power,
            self.min_heater_power,
            self.max_heater_power
        )

        # Store for next iteration
        self.prev_error = error.copy()

        # Log history
        self.error_history.append(np.mean(np.abs(error)))
        self.power_history.append(np.mean(self.heater_power))

        return self.heater_power

    def get_statistics(self):
        """Return control loop statistics"""
        return {
            'mean_error_pm': np.mean(self.error_history) * 1e12,
            'max_error_pm': np.max(self.error_history) * 1e12,
            'mean_power_mw': np.mean(self.power_history),
            'total_power_mw': np.sum(self.heater_power)
        }


class PhotonicChipThermalModel:
    """Simplified thermal model of photonic chip"""
    def __init__(self, num_rings, ambient_temp=300):
        self.num_rings = num_rings
        self.ambient_temp = ambient_temp

        # Thermal properties
        self.thermal_resistance = 5.0  # K/mW
        self.thermal_time_constant = 0.1  # seconds
        self.thermo_optic_coef = 80e-12  # pm/K

        # State
        self.temperature = np.ones(num_rings) * ambient_temp
        self.base_wavelength = 1550e-9

    def update(self, heater_powers, dt):
        """Update thermal state based on heater powers"""
        target_temp = self.ambient_temp + heater_powers * self.thermal_resistance

        # First-order thermal response
        alpha = dt / self.thermal_time_constant
        self.temperature += alpha * (target_temp - self.temperature)

        # Add some noise
        self.temperature += np.random.randn(self.num_rings) * 0.01

        return self.get_wavelengths()

    def get_wavelengths(self):
        """Calculate resonance wavelengths from temperature"""
        delta_T = self.temperature - self.ambient_temp
        delta_lambda = delta_T * self.thermo_optic_coef
        return self.base_wavelength + delta_lambda


# Simulation example
if __name__ == "__main__":
    import matplotlib.pyplot as plt

    num_rings = 64
    chip = PhotonicChipThermalModel(num_rings)
    controller = ThermalController(num_rings, target_wavelength=1550.5e-9)

    # Simulate control loop
    dt = 0.001  # 1 ms time step
    time_steps = 5000
    heater_powers = np.zeros(num_rings)

    wavelength_history = []
    error_history = []

    for t in range(time_steps):
        # Get current state
        wavelengths = chip.update(heater_powers, dt)

        # Update controller
        heater_powers = controller.update(wavelengths, dt)

        # Log
        wavelength_history.append(np.mean(wavelengths))
        error_history.append(np.mean(np.abs(wavelengths - 1550.5e-9)))

    # Plot results
    time = np.arange(time_steps) * dt * 1000  # ms

    plt.figure(figsize=(12, 4))
    plt.subplot(1, 2, 1)
    plt.plot(time, np.array(wavelength_history) * 1e9)
    plt.axhline(1550.5, color='r', linestyle='--', label='Target')
    plt.xlabel('Time (ms)')
    plt.ylabel('Wavelength (nm)')
    plt.legend()

    plt.subplot(1, 2, 2)
    plt.semilogy(time, np.array(error_history) * 1e12)
    plt.xlabel('Time (ms)')
    plt.ylabel('Error (pm)')
    plt.tight_layout()

    stats = controller.get_statistics()
    print(f"Final error: {stats['mean_error_pm']:.2f} pm")
    print(f"Total heater power: {stats['total_power_mw']:.1f} mW")