Quantizing Geospatial Neural Networks for Deployment

Quantizing geospatial neural networks converts 32-bit floating-point model weights into 8-bit integers, reducing memory footprint by roughly 75% and accelerating inference on CPU-bound edge devices or cloud raster pipelines. The technique works because spatial feature extraction relies on relative activation thresholds rather than absolute numerical precision. When a model learns to separate urban footprints from vegetated zones, those gradient boundaries remain stable after integer mapping. Static quantization in PyTorch achieves this through a calibration phase that records activation ranges across representative satellite or aerial tiles, followed by a kernel swap that replaces floating-point operations with optimized integer routines.

Implementation Workflow #

The following workflow demonstrates a complete, production-ready static quantization pipeline for a convolutional architecture processing multispectral rasters. The code assumes PyTorch 1.12+ and targets CPU inference via the fbgemm backend.

The static quantization steps are outlined below.

flowchart LR
    A["FP32 model<br/>(model.eval())"] --> B["Assign qconfig<br/>(fbgemm)"]
    B --> C["prepare()<br/>insert observers"]
    C --> D["Calibrate on<br/>representative tiles"]
    D --> E["convert()<br/>to INT8"]
    E --> F["Deploy<br/>(.pt / ONNX)"]

import torch
import torch.nn as nn
import torch.quantization

class GeoRasterCNN(nn.Module):
    def __init__(self, in_channels=3, num_classes=4):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, 16, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(16)
        self.relu = nn.ReLU()
        self.pool = nn.MaxPool2d(kernel_size=2)
        # Adjust spatial dimensions based on your tile size (e.g., 128x128 -> 64x64 after pooling)
        self.fc = nn.Linear(16 * 64 * 64, num_classes)

    def forward(self, x):
        x = self.pool(self.relu(self.bn1(self.conv1(x))))
        x = x.view(x.size(0), -1)
        return self.fc(x)

# 1. Initialize and freeze model
model = GeoRasterCNN(in_channels=3, num_classes=4)
model.eval()

# 2. Assign CPU quantization configuration
model.qconfig = torch.quantization.get_default_qconfig('fbgemm')

# 3. Insert observers into the graph
torch.quantization.prepare(model, inplace=True)

# 4. Calibration: pass representative geospatial tiles
# Ensure tiles match the exact preprocessing pipeline used in training
calibration_tiles = torch.randn(32, 3, 128, 128)
with torch.no_grad():
    _ = model(calibration_tiles)

# 5. Convert to quantized model
torch.quantization.convert(model, inplace=True)

# 6. Verify quantization
print(f"Quantized Conv1 weight dtype: {model.conv1.weight().dtype}")
print(f"Quantized FC weight dtype: {model.fc.weight().dtype}")

Calibration for Spatial Data #

Calibration determines the scale and zero-point parameters that map floating-point activations to 8-bit integers. Geospatial rasters exhibit high dynamic range due to varying illumination, atmospheric scattering, and sensor-specific reflectance scales. Observers must see data that matches your production pipeline exactly.

• Normalize before calibration: Quantization observers are scale-sensitive. Apply the same band-wise normalization (e.g., min-max scaling to [0, 1] or standardization) used during training before feeding tiles to the model.
• Use diverse tiles: Include scenes with cloud cover, shadows, and varying land cover. A homogeneous calibration set will produce narrow activation ranges, causing clipping during inference.
• Batch size matters: Run calibration with 32–128 tiles. Single-tile calibration fails to capture distribution tails and degrades spatial boundary detection.

Debugging Common Quantization Failures #

Deployment & Inference #

Once converted, the model runs entirely in integer arithmetic. Save it using torch.save(model.state_dict(), 'quantized_geo_model.pt') and load it with model.load_state_dict(torch.load(...)). For cross-platform deployment, export to ONNX using torch.onnx.export() and run it through ONNX Runtime, which natively supports quantized operators. When integrating with Python GIS workflows, pair the quantized model with rasterio or xarray for streaming tile ingestion, ensuring memory stays bounded during large-area inference.

For broader pipeline tuning, review strategies in Advanced Geospatial AI Optimization to align quantization with tiling strategies, batch scheduling, and hardware-specific kernels. The official PyTorch Static Quantization Documentation provides backend-specific configuration matrices and operator coverage tables.

Related Pages #

• Geospatial Machine Learning & AI