56 lines
2.5 KiB
Python
56 lines
2.5 KiB
Python
import kfp.dsl as dsl
|
|
|
|
@dsl.pipeline(
|
|
name='Prediction pipeline',
|
|
description='Generate slim models and optimize them with OpenVINO'
|
|
)
|
|
def tf_slim_optimize(
|
|
model_name = dsl.PipelineParam(name='model-name', value='resnet_v1_50'),
|
|
num_classes = dsl.PipelineParam(name='num-classes', value=1000),
|
|
checkpoint_url = dsl.PipelineParam(name='checkpoint-url', value='http://download.tensorflow.org/models/resnet_v1_50_2016_08_28.tar.gz'),
|
|
batch_size = dsl.PipelineParam(name='batch-size', value=1),
|
|
export_dir = dsl.PipelineParam(name='tf-export-dir', value='/tmp/export'),
|
|
generated_model_dir = dsl.PipelineParam(name='generated-model-dir', value='gs://your-bucket/folder'),
|
|
mo_options = dsl.PipelineParam(name='mo-option', value='--saved_model_dir .'),
|
|
input_numpy_file = dsl.PipelineParam(name='input-numpy-file', value='gs://intelai_public_models/images/imgs.npy'),
|
|
label_numpy_file = dsl.PipelineParam(name='label-numpy-file', value='gs://intelai_public_models/images/lbs.npy')):
|
|
|
|
slim = dsl.ContainerOp(
|
|
name='Create_model',
|
|
image='gcr.io/constant-cubist-173123/inference_server/ml_slim:6',
|
|
command=['python', 'slim_model.py'],
|
|
arguments=[
|
|
'--model_name', model_name,
|
|
'--batch_size', batch_size,
|
|
'--checkpoint_url', checkpoint_url,
|
|
'--num_classes', num_classes,
|
|
'--saved_model_dir', generated_model_dir,
|
|
'--export_dir', export_dir],
|
|
file_outputs={'generated-model-dir': '/tmp/saved_model_dir.txt'})
|
|
|
|
mo = dsl.ContainerOp(
|
|
name='Optimize_model',
|
|
image='gcr.io/constant-cubist-173123/inference_server/ml_mo:12',
|
|
command=['convert_model.py'],
|
|
arguments=[
|
|
'--input_path', '%s/saved_model.pb' % slim.output,
|
|
'--mo_options', mo_options,
|
|
'--output_path', slim.output],
|
|
file_outputs={'bin': '/tmp/bin_path.txt', 'xml': '/tmp/xml_path.txt'})
|
|
|
|
dsl.ContainerOp(
|
|
name='openvino-predict',
|
|
image='gcr.io/constant-cubist-173123/inference_server/ml_predict:6',
|
|
command=['python3', 'predict.py'],
|
|
arguments=[
|
|
'--model_bin', mo.outputs['bin'],
|
|
'--model_xml', mo.outputs['xml'],
|
|
'--input_numpy_file', input_numpy_file,
|
|
'--label_numpy_file', label_numpy_file,
|
|
'--output_folder', generated_model_dir],
|
|
file_outputs={})
|
|
|
|
if __name__ == '__main__':
|
|
import kfp.compiler as compiler
|
|
compiler.Compiler().compile(tf_slim_optimize, __file__ + '.tar.gz')
|