name: Train FC DNN using TF
description: Trains fully-connected neural network using Tensorflow
inputs:
  - {name: Transformed data dir,  type: GCSPath,  description: 'GCS path containing tf-transformed training and eval data.'} # type: {GCSPath: {path_type: Directory}}
  - {name: Schema,                type: GCSPath,       description: 'GCS json schema file path.'} # type: {GCSPath: {data_type: JSON}}
  - {name: Learning rate,         type: Float, default: '0.1',              description: 'Learning rate for training.'}
  - {name: Optimizer,             type: String, default: 'Adagrad', description: 'Optimizer for training. Valid values are: Adam, SGD, Adagrad. If not provided, tf.estimator default will be used.'}
  - {name: Hidden layer size,     type: String, default: '100',             description: 'Comma-separated hidden layer sizes. For example "200,100,50".'}
  - {name: Steps,                 type: Integer,                            description: 'Maximum number of training steps to perform. If unspecified, will honor epochs.'}
  #- {name: Epochs,                type: Integer, default: '',               description: 'Maximum number of training data epochs on which to train. If both "steps" and "epochs" are specified, the training job will run for "steps" or "epochs", whichever occurs first.'}
  - {name: Target,                type: String,                             description: 'Name of the column for prediction target.'}
  - {name: Preprocessing module,  type: GCSPath, default: '', description: 'GCS path to a python file defining "preprocess" and "get_feature_columns" functions.'} # type: {GCSPath: {data_type: Python}}
  - {name: Training output dir,   type: GCSPath,  description: 'GCS or local directory.'} # type: {GCSPath: {path_type: Directory}}
outputs:
  - {name: Training output dir,   type: GCSPath,  description: 'GCS or local directory.'} # type: {GCSPath: {path_type: Directory}}
  - {name: MLPipeline UI metadata, type: UI metadata}
implementation:
  container:
    image: gcr.io/ml-pipeline/ml-pipeline-kubeflow-tf-trainer:1.8.0-alpha.0
    command: [python2, -m, trainer.task]
    args: [
      --transformed-data-dir, {inputValue: Transformed data dir},
      --schema, {inputValue: Schema},
      --learning-rate, {inputValue: Learning rate},
      --optimizer, {inputValue: Optimizer},
      --hidden-layer-size, {inputValue: Hidden layer size},
      --steps, {inputValue: Steps},
#      --epochs, {inputValue: Epochs},
      --target, {inputValue: Target},
      --preprocessing-module, {inputValue: Preprocessing module},
      --job-dir, {inputValue: Training output dir},
      --exported-model-dir-uri-output-path, {outputPath: Training output dir},
      --ui-metadata-output-path, {outputPath: MLPipeline UI metadata},
    ]
