---
dataset:
    # Specify the name of the dataset
    name: mnist
    task:
        # This dataset is used to train for image classification,
        # so we specify the type of this task to be `mayo.task.image.Classify`.
        # In the future, we will support other tasks such as object detection # and NLP applications.
        type: mayo.task.image.Classify
        # Specify the number of class labels.
        num_classes: 10
        # Specify whether the class labels in the dataset contain a
        # background class, we assume that if a background class exists,
        # it is always the first label (zero-indexed).
        background_class: {has: false}
        # The shape of images in the dataset, `height` and `width`
        # are optional.
        shape:
            height: 28
            width: 28
            channels: 1
        # The preprocessing required by the training dataset as the standard
        # training pipeline for all neural networks using the same dataset.
        # Neural network models can further specify validation preprocessing
        # pipeline, as well as the final preprocessing stages regardless of
        # training/validation.
        preprocess:
            train: []
    # the paths pointing to the TFRecord files and the text file
    # containing class labels.
    path:
        train: mnist/train.tfrecord
        validate: mnist/test.tfrecord
        labels: mnist/labels.txt
    # number of examples in the training/validation datasets.
    num_examples_per_epoch:
        train: 60000
        validate: 10000

---
dataset:
    task:
        # We additionally specify that our LeNet-5 does not use a
        # background class.
        background_class: {use: false}
        preprocess:
            # We specify the shape of the input to our neural network.
            shape:
                height: 28
                width: 28
                channels: 1
            # We do not use additional preprocessing for validation.
            validate: null
            # For both validation or training, we add a final stage of
            # image preprocessing to transform the range of values from
            # [0, 1] to [-1, 1].
            final_cpu: {type: linear_map, scale: 2.0, shift: -1.0}
model:
    # Name of the model.
    name: lenet5
    # Layer definitions.
    layers:
        # `_init` is a partial layer definition, here we use `&init` as a
        # reference to the mapping, which contains an initializer for weights.
        _init: &init
            weights_initializer:
                type: tensorflow.truncated_normal_initializer
                stddev: 0.09
        # Definition for layer `conv0`
        conv0: &conv
            # It inherits the mapping referenced by `&init` above.
            <<: *init
            # The type of the layer is a convolution.
            type: convolution
            # We specify the size of the kernel, padding and number of
            # output channels.
            kernel_size: 5
            padding: valid
            num_outputs: 20
            # And in addition, a regularizer for the weights.
            weights_regularizer:
                type: tensorflow.contrib.layers.l2_regularizer
                scale: 0.004
        pool0: &pool
            # This defines a max pool layer with a 2x2 kernel, stride size 2,
            # and a valid padding.
            type: max_pool
            kernel_size: 2
            stride: 2
            padding: valid
        # `conv1` inherits all definitions in `conv0` as referenced by the
        # anchor `&conv`, and modifies the number of output channels to 50.
        conv1: {<<: *conv, num_outputs: 50}
        # `pool1` simply reuses the mapping referenced in `&pool`.
        pool1: *pool
        # all other layers should be straightforward.
        flatten: {type: flatten}
        dropout: {type: dropout, keep_prob: 0.5}
        fc1: &fc {<<: *init, type: fully_connected, num_outputs: 500}
        logits:
            <<: *fc
            # no activation function
            activation_fn: null
            # In Mayo, we allow `$(key_dot_path)` to access any value
            # pointed by the key path, here it substitutes the value with
            # the one in `dataset.task.num_classes`, which is 10 if we decide
            # to use the MNIST dataset.
            num_outputs: $(dataset.task.num_classes)
    graph:
        # A graph definition, stringing the layers specified above to form a
        # complete neural network.  Graph definition can also be a list of
        # such mappings, supporting diverging and converging paths.
        from: input
        with: [conv0, pool0, conv1, pool1, flatten, dropout, fc1, logits]
        to: output

---
# Mayo supports description importing using `_import`.
# Here it merges the trainer YAML with the contents of `exponential.yaml`
# in the same directory.
_import: exponential.yaml
train:
    learning_rate:
        # Here we specify the initial learning rate.
        _initial: 0.01
        # And the number of epochs required before we decay the learning rate.
        decay_steps: 300
    optimizer:
        # The type of optimizer used.
        type: tensorflow.train.GradientDescentOptimizer

---
train:
    learning_rate:
        # Here we use exponential decay as the policy to decay learning rate
        # after a certain number of epochs.
        type: tensorflow.train.exponential_decay
        # The default initial learning rate.
        _initial: 0.1
        # The default batch size which corresponds to the `_initial`
        # learning rate.
        _default_batch_size: 128
        # The factor used to decay the learning rate.
        decay_rate: 0.16
        # The actual initial learning rate, which scales correspondingly
        # to the current batch size used.
        learning_rate: !arith >
            $(train.learning_rate._initial) * math.sqrt(
                $(system.batch_size_per_gpu) * $(system.num_gpus) /
                $(train.learning_rate._default_batch_size))
        decay_steps: 30
        staircase: true