name: cnn

conda_env: conda_env.yaml
# Can have a docker_env instead of a conda_env, e.g.
# docker_env:
#    image:  mlflow-docker-example

entry_points:
  main:
    parameters:
      data_file: path
      regularization: {type: float, default: 0.1}
      batch_size: {type: int, default: 32}
      learning_rate: {type: float, default: 0.001}
      epochs: {type: int, default: 2}
    command: "python train_model.py with 'batch_size={batch_size}' 'learning_rate=${learning_rate}' 'epochs=${epochs}'"
  validate:
    parameters:
      data_file: path
      regularization: {type: float, default: 0.1}
      batch_size: {type: int, default: 32}
      learning_rate: {type: float, default: 0.001}
      epochs: {type: int, default: 2}
    command: "python train_model.py with 'batch_size={batch_size}' 'learning_rate=${learning_rate}' 'epochs=${epochs}'"