nvflare.edge.models.model module

class DeviceModel(net: Module)

Bases: Module

Model wrapper for classification with CrossEntropyLoss.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(input, label)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

export_model(net: Module, input_tensor_example, label_tensor_example)

Runs the full export pipeline to convert a PyTorch model into an Executorch-compatible format.

This includes:

• Capturing the forward graph

• Capturing backward graph for training (if applicable)

• Lowering to Edge dialect

• Lowering to Executorch format

Parameters:

• net (nn.Module) – The PyTorch model to export.

• input_tensor_example (torch.Tensor) – An example input tensor for tracing.

• label_tensor_example (torch.Tensor) – An example output/label tensor for training export.

Returns:

The final lowered and exported Executorch model.

Return type:

ExportedProgram

export_model_to_bytes(net: Module, input_shape, output_shape)

Exports a PyTorch model to ExecuTorch PTE format to be used in embedded or edge environments.

This function creates dummy input and label tensors based on the provided shapes, runs the model export pipeline (including lowering to Executorch), and returns the serialized model buffer.

Parameters:

• net (nn.Module) – The PyTorch model to export.

• input_shape (tuple) – The shape of the input tensor, e.g., (batch_size, channels, height, width).

• output_shape (tuple) – The shape of the output tensor, e.g., (batch_size, num_classes).

Returns:

The exported model (.pte) in bytes.