57 lines
2.2 KiB
Python
57 lines
2.2 KiB
Python
from typing import NamedTuple
|
|
from kfp.components import create_component_from_func
|
|
|
|
|
|
def automl_export_model_to_gcs(
|
|
model_path: str,
|
|
gcs_output_uri_prefix: str,
|
|
model_format: str = 'tf_saved_model',
|
|
) -> NamedTuple('Outputs', [
|
|
('model_directory', 'Uri'),
|
|
]):
|
|
"""Exports a trained model to a user specified Google Cloud Storage location.
|
|
|
|
Args:
|
|
model_path: The resource name of the model to export. Format: 'projects/<project>/locations/<location>/models/<model>'
|
|
gcs_output_uri_prefix: The Google Cloud Storage directory where the model should be written to. Must be in the same location as AutoML. Required location: us-central1.
|
|
model_format: The format in which the model must be exported. The available, and default, formats depend on the problem and model type. Possible formats: tf_saved_model, tf_js, tflite, core_ml, edgetpu_tflite. See https://cloud.google.com/automl/docs/reference/rest/v1/projects.locations.models/export?hl=en#modelexportoutputconfig
|
|
|
|
Annotations:
|
|
author: Alexey Volkov <alexey.volkov@ark-kun.com>
|
|
"""
|
|
from google.cloud import automl
|
|
|
|
client = automl.AutoMlClient()
|
|
response = client.export_model(
|
|
name=model_path,
|
|
output_config=automl.ModelExportOutputConfig(
|
|
model_format=model_format,
|
|
gcs_destination=automl.GcsDestination(
|
|
output_uri_prefix=gcs_output_uri_prefix,
|
|
),
|
|
),
|
|
)
|
|
|
|
print('Operation started:')
|
|
print(response.operation)
|
|
result = response.result()
|
|
metadata = response.metadata
|
|
print('Operation finished:')
|
|
print(metadata)
|
|
return (metadata.export_model_details.output_info.gcs_output_directory, )
|
|
|
|
|
|
if __name__ == '__main__':
|
|
automl_export_model_to_gcs_op = create_component_from_func(
|
|
automl_export_model_to_gcs,
|
|
output_component_file='component.yaml',
|
|
base_image='python:3.8',
|
|
packages_to_install=[
|
|
'google-cloud-automl==2.0.0',
|
|
],
|
|
annotations={
|
|
"author": "Alexey Volkov <alexey.volkov@ark-kun.com>",
|
|
"canonical_location": "https://raw.githubusercontent.com/Ark-kun/pipeline_components/master/components/gcp/automl/export_model_to_gcs/component.yaml",
|
|
},
|
|
)
|