tfm.core.export_base.ExportModule

Base Export Module.

tfm.core.export_base.ExportModule(
    params,
    model: Union[tf.Module, tf.keras.Model],
    inference_step: Optional[Callable[..., Any]] = None,
    *,
    preprocessor: Optional[Callable[..., Any]] = None,
    postprocessor: Optional[Callable[..., Any]] = None
)

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

@abc.abstractmethod
get_inference_signatures(
    function_keys: Dict[Text, Text]
) -> Mapping[Text, Any]

Get defined function signatures.

serve

View source

@abc.abstractmethod
serve() -> Mapping[Text, tf.Tensor]

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.