Legacy model weights import

[talmo]

Core SLEAP models are saved as Keras HDF5 files, which makes them fairly easily readable without actually depending on TensorFlow or Keras.

The task would be to understand the structure of those HDF5 files, load up the weights, possibly transpose the ordering to match the kernel ordering in PyTorch, and assign those values to the appropriate nn.modules.

Legacy models with weights artifacts for testing can be found here: https://github.com/talmolab/sleap/tree/develop/tests/data/models

See also:

• Legacy config support #140 which is a pre-requisite so that we can instantiate the models as PyTorch nn.modules in the first place.
• Map legacy SLEAP json configs to SLEAP-NN OmegaConf objects #162 which implements legacy config loading.

[talmo]

From Claude Code for reference:

Legacy Model Weights Import - Analysis and Conversion Plan

I've completed an initial analysis of the legacy Keras HDF5 model files and developed a conversion plan for importing them into the PyTorch-based sleap-nn framework.

Key Findings

1. Weight Format Differences

• Keras stores conv weights as (H, W, C_in, C_out)
• PyTorch expects conv weights as (C_out, C_in, H, W)
• This requires transposing weights during conversion

2. Layer Structure in HDF5 Files

# Example from centroid model:
model_weights/
├── stack0_enc0_conv0/stack0_enc0_conv0/kernel:0  # Shape: (3, 3, 1, 16)
├── stack0_enc0_conv0/stack0_enc0_conv0/bias:0    # Shape: (16,)
├── stack0_enc1_conv0/stack0_enc1_conv0/kernel:0  # Shape: (3, 3, 16, 24)
├── ...
├── CentroidConfmapsHead_0/CentroidConfmapsHead_0/kernel:0  # Shape: (1, 1, 36, 1)
└── OffsetRefinementHead_0/OffsetRefinementHead_0/kernel:0  # Shape: (1, 1, 36, 2)

3. Layer Naming Patterns

• Legacy: stack0_enc{block}_conv{conv_idx} for encoder layers
• Legacy: stack0_dec{block}_s{stride1}_to_s{stride2}_* for decoder layers
• PyTorch: enc.encoder_stack[{block}].blocks.{idx} (includes pooling/activation layers)

Proposed Implementation Approach

Step 1: Weight Conversion Utilities

def convert_keras_to_pytorch_weights(keras_weight: np.ndarray, layer_type: str) -> torch.Tensor:
    """Convert Keras weight format to PyTorch format."""
    if layer_type == "conv2d":
        # Keras: (H, W, C_in, C_out) -> PyTorch: (C_out, C_in, H, W)
        return torch.from_numpy(keras_weight.transpose(3, 2, 0, 1))
    elif layer_type == "bias":
        return torch.from_numpy(keras_weight)

Step 2: Layer Name Mapping

• Parse legacy layer names to extract block/conv indices
• Map to corresponding PyTorch module paths
• Handle special cases (middle blocks, heads)

Step 3: Model Architecture Matching

• Load training_job_config.json to get architecture params
• Create PyTorch model with matching configuration
• Map legacy parameters to PyTorch equivalents

Implementation Priority

1. Centroid model (simplest architecture)
2. Centered instance model (similar structure)
3. Single instance model (different input channels)
4. Bottom-up model (most complex with PAF heads)