pytorch_mnist.yaml

apiVersion: "kubeflow.org/v1alpha1" 
kind: StudyJob

metadata:
  namespace: kubeflow
  labels:
    controller-tools.k8s.io: "1.0"
  name: pytorchmnist
spec:
  studyName: pytorchmnist
  owner: crd
  optimizationtype: maximize
  objectivevaluename: Validation-accuracy
  optimizationgoal: 0.99
  requestcount: 4
  metricsnames:
    - accuracy
  parameterconfigs:
    - name: --lr
      parametertype: double
      feasible:
        min: "0.01"
        max: "0.03"
    - name: --num-layers
      parametertype: int
      feasible:
        min: "2"
        max: "5"
    - name: --optimizer
      parametertype: categorical
      feasible:
        list:
        - sgd
        - adam
        - ftrl
  workerSpec:
    goTemplate:
        rawTemplate: |-
          apiVersion: batch/v1
          kind: Job
          metadata:
            name: {{.WorkerID}}
            namespace: kubeflow
          spec:
            template:
              spec:
                containers:
                - name: pytorch
                  image: gcr.io/kubeflow-ci/pytorch-dist-mnist-test:v1.0
                  command:
                  - "python"
                  - "pytorch_mnist/train_mnist.py"
                  {{- with .HyperParameters}}
                  {{- range .}}
                  - "{{.Name}}={{.Value}}"
                  {{- end}}
                  {{- end}}
                restartPolicy: Never
  suggestionSpec:
    suggestionAlgorithm: "random"
    requestNumber: 3