Port TensorFlow Quickstart to NNI

This is a modified version of TensorFlow quickstart.

It can be run directly and will have the exact same result as original version.

Furthermore, it enables the ability of auto tuning with an NNI experiment, which will be detailed later.

It is recommended to run this script directly first to verify the environment.

There are 3 key differences from the original version:

  1. In Get optimized hyperparameters part, it receives generated hyperparameters.

  2. In (Optional) Report intermediate results part, it reports per-epoch accuracy metrics.

  3. In Report final result part, it reports final accuracy.

import nni
import tensorflow as tf

Hyperparameters to be tuned

These are the hyperparameters that will be tuned later.

params = {
    'dense_units': 128,
    'activation_type': 'relu',
    'dropout_rate': 0.2,
    'learning_rate': 0.001,
}

Get optimized hyperparameters

If run directly, nni.get_next_parameter() is a no-op and returns an empty dict. But with an NNI experiment, it will receive optimized hyperparameters from tuning algorithm.

optimized_params = nni.get_next_parameter()
params.update(optimized_params)
print(params)

Out:

{'dense_units': 128, 'activation_type': 'relu', 'dropout_rate': 0.2, 'learning_rate': 0.001}

Load dataset

mnist = tf.keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

Build model with hyperparameters

model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(params['dense_units'], activation=params['activation_type']),
    tf.keras.layers.Dropout(params['dropout_rate']),
    tf.keras.layers.Dense(10)
])

adam = tf.keras.optimizers.Adam(learning_rate=params['learning_rate'])
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
model.compile(optimizer=adam, loss=loss_fn, metrics=['accuracy'])

(Optional) Report intermediate results

The callback reports per-epoch accuracy to show learning curve in the web portal. You can also leverage the metrics for early stopping with NNI assessors.

This part can be safely skipped and the experiment will work fine.

callback = tf.keras.callbacks.LambdaCallback(
    on_epoch_end = lambda epoch, logs: nni.report_intermediate_result(logs['accuracy'])
)

Train and evluate the model

model.fit(x_train, y_train, epochs=5, verbose=2, callbacks=[callback])
loss, accuracy = model.evaluate(x_test, y_test, verbose=2)

Out:

Epoch 1/5
[2022-03-21 01:25:00] INFO (nni/MainThread) Intermediate result: 0.9153500199317932  (Index 0)
1875/1875 - 17s - loss: 0.2914 - accuracy: 0.9154 - 17s/epoch - 9ms/step
Epoch 2/5
[2022-03-21 01:25:18] INFO (nni/MainThread) Intermediate result: 0.9588666558265686  (Index 1)
1875/1875 - 18s - loss: 0.1387 - accuracy: 0.9589 - 18s/epoch - 10ms/step
Epoch 3/5
[2022-03-21 01:25:38] INFO (nni/MainThread) Intermediate result: 0.9677000045776367  (Index 2)
1875/1875 - 20s - loss: 0.1073 - accuracy: 0.9677 - 20s/epoch - 11ms/step
Epoch 4/5
[2022-03-21 01:25:56] INFO (nni/MainThread) Intermediate result: 0.9738666415214539  (Index 3)
1875/1875 - 18s - loss: 0.0866 - accuracy: 0.9739 - 18s/epoch - 10ms/step
Epoch 5/5
[2022-03-21 01:26:16] INFO (nni/MainThread) Intermediate result: 0.977483332157135  (Index 4)
1875/1875 - 21s - loss: 0.0728 - accuracy: 0.9775 - 21s/epoch - 11ms/step
313/313 - 2s - loss: 0.0702 - accuracy: 0.9776 - 2s/epoch - 6ms/step

Report final result

Report final accuracy to NNI so the tuning algorithm can suggest better hyperparameters.

nni.report_final_result(accuracy)

Out:

[2022-03-21 01:27:08] INFO (nni/MainThread) Final result: 0.9775999784469604

Total running time of the script: ( 2 minutes 27.156 seconds)

Gallery generated by Sphinx-Gallery