torch-layer-segregator 0.1.0

Tag PyTorch submodules and export them as Torch, TorchScript, or Core ML models.

torch_layer_segregator

Tag PyTorch submodules and export them as standalone Torch, TorchScript, or Core ML models.

Useful for profiling individual layers, benchmarking conversions, and building per-layer deployment pipelines without manually tracing each submodule.

Install

pip install torch_layer_segregator

Requirements: Python 3.10+, PyTorch 2.7+. Core ML export also requires coremltools 9.0+ on macOS.

Quick start

import torch.nn as nn
from torch_layer_segregator import TorchLayerSaver, make_saveable_layer

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.block = make_saveable_layer(nn.Conv2d(3, 16, 3, padding=1), key="conv_block")

    def forward(self, x):
        return self.block(x)

model = MyModel()
input_shape = (1, 3, 224, 224)

saver = TorchLayerSaver(model, "exports/my_model", input_shape, device="cpu")
saver.save_torch_model_layers()
saver.save_script_model_layers()
saver.save_coreml_model_layers()

How it works

1. Tag layers — wrap submodules with make_saveable_layer to mark them for export.
2. Infer shapes — TorchLayerSaver runs one forward pass and records each tagged layer's input and output shapes.
3. Export — save tagged layers to disk in the format you need.

Tagged layers are written under subdirectories of your output path, each with a layer_details.json manifest.

Method	Directory	Output
save_torch_model_layers()	torch_layers/	.pt (metadata + state_dict)
save_script_model_layers()	script_layers/	TorchScript .ts
save_coreml_model_layers()	coreml_layers/	Core ML .mlpackage

If no key is passed to make_saveable_layer, the filename is derived from the module's dotted name (e.g. encoder.block0 → encoder_block0).

Core ML conversion traces each layer and exports an ML Program with float16 precision for all compute units (iOS 26+ deployment target).

API

make_saveable_layer(module, key=None, input_shape=None)
Mark a nn.Module for export. Returns the same module for inline use in __init__.

TorchLayerSaver(model, out_dir, input_shape, device="cpu")
Run shape inference and coordinate exports. input_shape is the full model input tuple (batch included). device is used for inference and scripted/Core ML export.

All save_* methods accept an optional out_dir to override the default subdirectory.