This document provides a detailed overview of the SensorAugmentor framework architecture, explaining the core design principles, components, data flow, and implementation details.
- Design Philosophy
- System Architecture
- Core Components
- Data Flow
- Model Architecture
- Training Pipeline
- Extension Points
- Performance Considerations
- Future Architecture Directions
The SensorAugmentor framework was built on several key design principles:
-
Signal Enhancement through Deep Learning: Leveraging neural networks to enhance low-quality sensor data by learning the mapping to high-quality signals.
-
Multi-Objective Optimization: Simultaneously optimizing for signal reconstruction and actuator command prediction through shared latent representations.
-
Modularity and Extensibility: Component-based design allowing users to replace or extend specific parts without affecting the entire system.
-
Minimalism: Preferring simple, focused components that do one thing well over complex, monolithic implementations.
-
Production Readiness: Designed with real-world deployment scenarios in mind, including serialization, testing, and performance optimization.
The SensorAugmentor framework consists of the following high-level layers:
- Data Layer: Handles data loading, preprocessing, augmentation, and batching.
- Model Layer: Implements the neural network architecture for signal processing.
- Training Layer: Manages the training process, optimization, and evaluation.
- Serialization Layer: Handles model saving, loading, and versioning.
- Deployment Layer: Facilitates model deployment in various environments.
The following diagram illustrates the high-level architecture:
+-------------------+ +-------------------+ +-------------------+
| Data Layer | | Model Layer | | Training Layer |
+-------------------+ +-------------------+ +-------------------+
| - Data Loading | | - Encoder | | - Loss Functions |
| - Preprocessing |---->| - Latent Space |---->| - Optimizers |
| - Augmentation | | - Decoders | | - Validation |
| - Batching | | - ResidualBlocks | | - Early Stopping |
+-------------------+ +-------------------+ +-------------------+
| | |
v v v
+-------------------+ +-------------------+ +-------------------+
| Serialization | | Deployment | | Inference |
| Layer | | Layer | | Layer |
+-------------------+ +-------------------+ +-------------------+
| - Model Saving | | - API Services | | - Batch Inference |
| - Model Loading | | - Containerization| | - Real-time |
| - Versioning | | - Edge Deployment | | Processing |
+-------------------+ +-------------------+ +-------------------+
The ResidualBlock is a fundamental building block that implements skip connections to help with gradient flow during training and enable deeper networks:
class ResidualBlock(nn.Module):
def __init__(self, dim):
super(ResidualBlock, self).__init__()
self.block = nn.Sequential(
nn.Linear(dim, dim),
nn.ReLU(),
nn.Linear(dim, dim)
)
def forward(self, x):
return x + self.block(x)Key features:
- Simple yet effective design with two linear layers and a ReLU activation
- Skip connection that adds input directly to the output
- Preserves input dimension, making it easy to stack multiple blocks
The SensorAugmentor class is the main model that implements the sensor data enhancement and actuator command prediction:
class SensorAugmentor(nn.Module):
def __init__(self, sensor_dim, hidden_dim=64, output_dim=None, num_resblocks=2):
super(SensorAugmentor, self).__init__()
self.sensor_dim = sensor_dim
self.hidden_dim = hidden_dim
self.output_dim = output_dim if output_dim is not None else sensor_dim
# Encoder (Low-quality sensor to latent space)
self.encoder = nn.Sequential(
nn.Linear(sensor_dim, hidden_dim),
nn.ReLU()
)
# Post-encoding residual blocks
self.post_encoding_resblock = self._build_resblocks(hidden_dim, num_resblocks)
# High-quality signal reconstructor
self.hq_reconstructor = nn.Linear(hidden_dim, sensor_dim)
# Actuator command head (if output_dim is provided)
self.actuator_head = nn.Linear(hidden_dim, self.output_dim)
def _build_resblocks(self, dim, num_blocks):
blocks = []
for _ in range(num_blocks):
blocks.append(ResidualBlock(dim))
return nn.Sequential(*blocks)
def forward(self, x_lq, x_hq=None):
# Encode low-quality input
encoded_lq = self.encoder(x_lq)
encoded_lq = self.post_encoding_resblock(encoded_lq)
# Reconstruct high-quality signal
reconstructed_hq = self.hq_reconstructor(encoded_lq)
# Generate actuator command
actuator_command = self.actuator_head(encoded_lq)
# Encode high-quality signal if provided (for training)
encoded_hq = None
if x_hq is not None:
encoded_hq = self.encoder(x_hq)
encoded_hq = self.post_encoding_resblock(encoded_hq)
return reconstructed_hq, actuator_command, encoded_lq, encoded_hqKey features:
- Encoder that maps sensor data to a latent representation
- ResidualBlocks for improved gradient flow
- Separate heads for reconstructing high-quality signals and predicting actuator commands
- Optional encoding of high-quality signals for training purposes
The SyntheticSensorDataset class provides synthetic data generation for training and testing:
class SyntheticSensorDataset(Dataset):
def __init__(self, num_samples=1000, sensor_dim=10, output_dim=None, noise_level=0.2, seed=None):
self.num_samples = num_samples
self.sensor_dim = sensor_dim
self.output_dim = output_dim if output_dim is not None else sensor_dim
self.noise_level = noise_level
if seed is not None:
torch.manual_seed(seed)
np.random.seed(seed)
# Generate "high-quality" sensor data (ground truth)
self.x_hq = torch.randn(num_samples, sensor_dim)
# Generate "low-quality" sensor data (noisy version)
self.x_lq = self.x_hq + noise_level * torch.randn(num_samples, sensor_dim)
# Generate "actuator commands" (simulated relationship)
self.y_cmd = torch.sin(self.x_hq[:, 0:1]) + torch.cos(self.x_hq[:, 1:2])
if self.output_dim > 1:
# Add more simulated relationships for multi-dimensional output
extra_dims = torch.zeros(num_samples, self.output_dim - 1)
for i in range(1, self.output_dim):
if i < sensor_dim:
extra_dims[:, i-1] = torch.tanh(self.x_hq[:, i])
self.y_cmd = torch.cat([self.y_cmd, extra_dims], dim=1)
def __len__(self):
return self.num_samples
def __getitem__(self, idx):
return self.x_lq[idx], self.x_hq[idx], self.y_cmd[idx]Key features:
- Customizable number of samples, dimensions, and noise levels
- Deterministic generation using optional seed for reproducibility
- Synthetic relationships between high-quality signals and actuator commands
The EarlyStopper class implements early stopping to prevent overfitting:
class EarlyStopper:
def __init__(self, patience=5, min_delta=0):
self.patience = patience
self.min_delta = min_delta
self.counter = 0
self.min_val_loss = float('inf')
def early_stop(self, val_loss):
if val_loss < self.min_val_loss - self.min_delta:
self.min_val_loss = val_loss
self.counter = 0
else:
self.counter += 1
if self.counter >= self.patience:
return True
return FalseKey features:
- Configurable patience and minimum improvement threshold
- Simple interface for tracking validation loss
- Automatic determination of when to stop training
The SensorAugmentor framework processes data through the following flow:
- Data Ingestion: Raw sensor data is loaded and parsed.
- Preprocessing: Data is cleaned, normalized, and prepared for model input.
- Model Processing:
- Low-quality sensor data is encoded into a latent representation.
- The latent representation is processed through residual blocks.
- The processed representation is used to reconstruct high-quality signals and predict actuator commands.
- Postprocessing: Model outputs are denormalized and prepared for use.
- Feedback Loop: During training, performance metrics are used to update model parameters.
The following diagram illustrates the data flow during training:
+-------------------------+
| Low-Quality Sensor Data |
+-------------------------+
|
v
+----------------+
| Encoder |
+----------------+
|
v
+----------------+
| Residual Blocks |
+----------------+
|
v
+-------------------+-------------------+
| |
v v
+---------------------+ +---------------------+
| HQ Reconstructor | | Actuator Head |
+---------------------+ +---------------------+
| |
v v
+---------------------+ +---------------------+
| Reconstructed HQ | | Actuator Commands |
| Signal | | |
+---------------------+ +---------------------+
| |
v v
+---------------------+ +---------------------+
| Compare with | | Compare with |
| True HQ Signal | | True Commands |
+---------------------+ +---------------------+
| |
+-------------------+-------------------+
|
v
+-------------------+
| Combined Loss |
+-------------------+
|
v
+-------------------+
| Backpropagation |
+-------------------+
|
v
+-------------------+
| Parameter Update |
+-------------------+
During inference, the flow is simplified:
+-------------------------+
| Low-Quality Sensor Data |
+-------------------------+
|
v
+----------------+
| Encoder |
+----------------+
|
v
+----------------+
| Residual Blocks |
+----------------+
|
v
+-------------------+-------------------+
| |
v v
+---------------------+ +---------------------+
| HQ Reconstructor | | Actuator Head |
+---------------------+ +---------------------+
| |
v v
+---------------------+ +---------------------+
| Reconstructed HQ | | Actuator Commands |
| Signal | | |
+---------------------+ +---------------------+
The SensorAugmentor model architecture is based on a shared encoder with multiple output heads. This design allows the model to learn a common latent representation that captures important features of the sensor data, which can then be used for multiple tasks.
The encoder consists of:
- An initial linear layer that projects the input sensor data into a higher-dimensional latent space
- A non-linear activation function (ReLU) to introduce non-linearity
- A series of residual blocks to process the latent representation further
The encoder learns to map noisy, low-quality sensor data to a cleaner, more informative latent representation. By training on pairs of low-quality and high-quality data, the encoder learns to extract meaningful features from noisy inputs.
Residual blocks are key to the model's performance, allowing for deeper networks without the vanishing gradient problem. Each block includes:
- A skip connection that preserves the input
- A non-linear transformation path with two linear layers and a ReLU activation
- An additive combination of the skip connection and transformed path
The residual connections help gradient flow during backpropagation, enabling more effective training of deeper networks.
The model has two output heads:
- HQ Reconstructor: Reconstructs high-quality sensor signals from the latent representation
- Actuator Head: Predicts actuator commands from the same latent representation
This multi-task design forces the model to learn a general, informative latent representation that captures relevant features for both tasks.
The training uses a composite loss function that combines:
- Reconstruction Loss: Mean squared error between reconstructed and true high-quality signals
- Encoding Loss: Mean squared error between encodings of low-quality and high-quality signals
- Actuator Loss: Mean squared error between predicted and true actuator commands
def calculate_loss(reconstructed_hq, act_command, encoded_lq, encoded_hq, x_hq, y_cmd):
loss_recon = F.mse_loss(reconstructed_hq, x_hq)
loss_encoding = F.mse_loss(encoded_lq, encoded_hq)
loss_act = F.mse_loss(act_command, y_cmd)
# Combined loss with weights
total_loss = loss_recon + 0.1 * loss_encoding + 0.5 * loss_act
return total_loss, loss_recon, loss_encoding, loss_actThe weights for each component of the loss can be adjusted based on the specific use case.
The training pipeline consists of several key components:
- Dataset Creation: Create or load datasets for training and validation
- Data Loaders: Create PyTorch data loaders for efficient batching and shuffling
- Normalization: Apply normalization to standardize input data
The core training loop:
def train_model(model, train_loader, val_loader, num_epochs=100, lr=1e-3, patience=5):
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
early_stopper = EarlyStopper(patience=patience)
train_losses = []
val_losses = []
for epoch in range(num_epochs):
# Training phase
model.train()
train_loss = 0.0
for x_lq, x_hq, y_cmd in train_loader:
optimizer.zero_grad()
# Forward pass
reconstructed_hq, act_command, encoded_lq, encoded_hq = model(x_lq, x_hq)
# Calculate loss
loss, _, _, _ = calculate_loss(reconstructed_hq, act_command, encoded_lq, encoded_hq, x_hq, y_cmd)
# Backward pass and optimize
loss.backward()
optimizer.step()
train_loss += loss.item()
train_loss /= len(train_loader)
train_losses.append(train_loss)
# Validation phase
model.eval()
val_loss = 0.0
with torch.no_grad():
for x_lq, x_hq, y_cmd in val_loader:
reconstructed_hq, act_command, encoded_lq, encoded_hq = model(x_lq, x_hq)
loss, _, _, _ = calculate_loss(reconstructed_hq, act_command, encoded_lq, encoded_hq, x_hq, y_cmd)
val_loss += loss.item()
val_loss /= len(val_loader)
val_losses.append(val_loss)
print(f"Epoch {epoch+1}/{num_epochs}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}")
# Early stopping
if early_stopper.early_stop(val_loss):
print(f"Early stopping at epoch {epoch+1}")
break
return train_losses, val_lossesAfter training, the model is evaluated on test data to measure:
- Signal Reconstruction Quality: How well the model reconstructs high-quality signals
- Actuator Command Accuracy: How accurately the model predicts actuator commands
- Latent Space Quality: How informative and structured the learned latent space is
The SensorAugmentor framework is designed to be modular and extensible. Key extension points include:
Users can replace the default encoder with custom architectures:
class CNNEncoder(nn.Module):
def __init__(self, sensor_dim, hidden_dim):
super(CNNEncoder, self).__init__()
self.conv1 = nn.Conv1d(1, 16, kernel_size=3, padding=1)
self.conv2 = nn.Conv1d(16, 32, kernel_size=3, padding=1)
self.pool = nn.MaxPool1d(2)
self.flatten = nn.Flatten()
self.fc = nn.Linear(32 * (sensor_dim // 2), hidden_dim)
self.relu = nn.ReLU()
def forward(self, x):
# Reshape for 1D convolution [batch, 1, sensor_dim]
x = x.unsqueeze(1)
x = self.relu(self.conv1(x))
x = self.pool(x)
x = self.relu(self.conv2(x))
x = self.flatten(x)
x = self.relu(self.fc(x))
return x
# Create custom SensorAugmentor with CNN encoder
class CNNSensorAugmentor(SensorAugmentor):
def __init__(self, sensor_dim, hidden_dim=64, output_dim=None, num_resblocks=2):
super(CNNSensorAugmentor, self).__init__(sensor_dim, hidden_dim, output_dim, num_resblocks)
# Replace the default encoder with CNN encoder
self.encoder = CNNEncoder(sensor_dim, hidden_dim)Users can implement alternative residual block designs:
class GatedResidualBlock(nn.Module):
def __init__(self, dim):
super(GatedResidualBlock, self).__init__()
self.block = nn.Sequential(
nn.Linear(dim, dim),
nn.ReLU(),
nn.Linear(dim, dim)
)
self.gate = nn.Sequential(
nn.Linear(dim, dim),
nn.Sigmoid()
)
def forward(self, x):
block_output = self.block(x)
gate_output = self.gate(x)
return x + gate_output * block_outputUsers can define custom loss functions for specific applications:
def custom_loss(reconstructed_hq, act_command, encoded_lq, encoded_hq, x_hq, y_cmd):
# Standard losses
loss_recon = F.mse_loss(reconstructed_hq, x_hq)
loss_encoding = F.mse_loss(encoded_lq, encoded_hq)
loss_act = F.mse_loss(act_command, y_cmd)
# Add custom frequency-domain loss for sensor data
freq_recon = torch.rfft(reconstructed_hq, 1)
freq_hq = torch.rfft(x_hq, 1)
loss_freq = F.mse_loss(freq_recon, freq_hq)
# Combined loss with weights
total_loss = loss_recon + 0.1 * loss_encoding + 0.5 * loss_act + 0.3 * loss_freq
return total_lossSeveral optimizations improve the performance of SensorAugmentor:
- Batch Processing: All operations support batch processing for efficient parallel computation
- Minimal Architecture: The base model uses a minimal number of layers to reduce computation
- Shared Encoder: Using a shared encoder for both tasks reduces computational overhead
- Parameter Sharing: Shared encoder reduces the total number of parameters
- Residual Connections: Allow deeper networks with fewer total parameters
- Linear Layers: Using linear layers instead of convolutional layers for the base model reduces memory requirements
- Early Stopping: Prevents unnecessary training iterations
- Composite Loss: Multi-objective loss function allows simultaneous optimization of all tasks
- Residual Connections: Improve gradient flow for faster convergence
Future versions of the SensorAugmentor framework may include:
-
Advanced Model Architectures:
- Transformer-based encoders for sequence modeling
- Graph neural networks for spatial sensor arrays
- Variational autoencoders for probabilistic latent spaces
-
Enhanced Training Methods:
- Self-supervised pretraining on unlabeled sensor data
- Meta-learning for quick adaptation to new sensor types
- Continual learning for adapting to sensor drift over time
-
Deployment Optimizations:
- Model quantization for edge deployment
- Model pruning for reduced size
- JIT compilation for optimized inference
-
Uncertainty Quantification:
- Bayesian neural networks for uncertainty estimation
- Ensemble methods for robust predictions
- Confidence calibration for reliable decision-making
-
Multi-Modal Integration:
- Fusion of heterogeneous sensor types (e.g., visual, acoustic, vibration)
- Cross-modal learning for improved signal reconstruction
- Multi-sensor synchronization and alignment