examples/demos/simple_pipeline/gpu-example-pipeline.py

47 lines
1.4 KiB
Python
Executable File

#!/usr/bin/env python3
import kfp.dsl as kfp
def training_op(learning_rate: float,
num_layers: int,
optimizer='ftrl',
step_name='training'):
return kfp.ContainerOp(
name=step_name,
image='katib/mxnet-mnist-example',
command=['python', '/mxnet/example/image-classification/train_mnist.py'],
arguments=[
'--batch-size', '64',
'--lr', learning_rate,
'--num-layers', num_layers,
'--optimizer', optimizer
],
file_outputs={'output': '/etc/timezone'}
)
def postprocessing_op(output,
step_name='postprocessing'):
return kfp.ContainerOp(
name=step_name,
image='library/bash:4.4.23',
command=['sh', '-c'],
arguments=['echo "%s"' % output]
)
@kfp.pipeline(
name='Pipeline GPU Example',
description='Demonstrate the Kubeflow pipelines SDK with GPUs'
)
def kubeflow_training(
learning_rate: kfp.PipelineParam = kfp.PipelineParam(name='learningrate', value=0.1),
num_layers: kfp.PipelineParam = kfp.PipelineParam(name='numlayers', value='2'),
optimizer: kfp.PipelineParam = kfp.PipelineParam(name='optimizer', value='ftrl')):
training = training_op(learning_rate, num_layers, optimizer).set_gpu_limit(1)
postprocessing = postprocessing_op(training.output) # pylint: disable=unused-variable
if __name__ == '__main__':
import kfp.compiler as compiler
compiler.Compiler().compile(kubeflow_training, __file__ + '.tar.gz')