tfm.core.export_base.ExportModule

Base Export Module.

params A dataclass for parameters to the module.
model A model instance which contains weights and forward computation.
inference_step An optional callable to forward-pass the model. If not specified, it creates a parital function with model as an required kwarg.
preprocessor An optional callable to preprocess the inputs.
postprocessor An optional callable to postprocess the model outputs.

Methods

get_inference_signatures

View source

Get defined function signatures.

serve

View source

The bare inference function which should run on all devices.

Expecting tensors are passed in through keyword arguments. Returns a dictionary of tensors, when the keys will be used inside the SignatureDef.