mirror of https://github.com/kubeflow/examples.git
47 lines
1.4 KiB
Python
Executable File
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')
|