Maggy Distributed Training on MNIST Dataset

Maggy enables you to train with Tensorflow distributed optimizers. Using Maggy, you have to make minimal changes in train your model in a distributed fashion.

0. Spark Session

Make sure you have a running Spark Session/Context available. On Hopsworks, just run your notebook to start the spark application.

from maggy import experiment

---------------------------------------------------------------------------

ModuleNotFoundError                       Traceback (most recent call last)

<ipython-input-1-308631a69d5b> in <module>
----> 1 from maggy import experiment


/srv/hops/anaconda/envs/theenv/lib/python3.7/site-packages/maggy/experiment.py in <module>
     29 from typing import Callable
     30 
---> 31 from maggy import util
     32 from maggy.core.environment.singleton import EnvSing
     33 from maggy.experiment_config.lagom import LagomConfig


/srv/hops/anaconda/envs/theenv/lib/python3.7/site-packages/maggy/util.py in <module>
     22 
     23 import numpy as np
---> 24 from pyspark import TaskContext
     25 from pyspark.sql import SparkSession
     26 


ModuleNotFoundError: No module named 'pyspark'

1. Model definition

Let’s define the model we want to train. The layers of the model have to be defined in the __init__ function.

Do not instantiate the class, otherwise you won’t be able to use Maggy.

# Model definition.
import tensorflow as tf
from tensorflow import keras
    
class NeuralNetwork(tf.keras.Model):

    def __init__(self, nlayers):
        super().__init__()
        self.conv1 = keras.layers.Conv2D(28, 2, activation='relu')
        self.flatten = keras.layers.Flatten()
        self.d1 = keras.layers.Dense(32, activation='relu')
        self.d2 = keras.layers.Dense(10, activation='softmax')
        
    def call(self, x):
        x = self.conv1(x)
        x = self.flatten(x)
        x = self.d1(x)
        return self.d2(x)

model = NeuralNetwork

2. Dataset creation

You can create the dataset here and pass it to the TfDistributedConfig, or creating it in the training function.

In this example, we are downloading the dataset from tensorflow.

import numpy as np

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
x_train = np.reshape(x_train, (60000, 28, 28, 1))
x_test = np.reshape(x_test, (10000, 28, 28, 1))

3. Defining the training function

The programming model is that you wrap the code containing the model training inside a wrapper function. Inside that wrapper function provide all imports and parts that make up your experiment.

The function should return the metric that you want to optimize for. This should coincide with the metric being reported in the Keras callback (see next point). You can return the metric list, in this case only the loss element will be printed.

def training_function(model, train_set, test_set, hparams):

    from tensorflow import keras
    
    # Define training parameters
    num_epochs = 10
    batch_size = 256
    learning_rate = 0.1

    criterion = keras.losses.SparseCategoricalCrossentropy()
    optimizer = keras.optimizers.SGD(learning_rate=learning_rate,momentum=0.9,decay=1e-5)
    
    model = model(nlayers = 2)
    
    model.compile(optimizer=optimizer, loss=criterion, metrics=["accuracy"])
    
    model.fit(train_set,
              batch_size=batch_size,
              epochs=num_epochs,
              )

    print("Testing")
    
    loss = model.evaluate(
        test_set,
        batch_size=32)
    
    return loss

4. Configuring the experiment

In order to use maggy distributed training, we have to configure the training model, we can pass it to TfDistributedConfig. the model class has to be an implementation of tf.keras.Model. We can also define train_set, test_set and eventually the model_parameters. model_parameters is a dictionary containing the parameters to be used in the __init__ function of your model.

from maggy.experiment_config.tf_distributed import TfDistributedConfig

#define the constructor parameters of your model
model_parameters = {
    'train_batch_size': 30000,
    'test_batch_size': 5000,
    'nlayers': 2
}

#pass the model parameters in the last 
config = TfDistributedConfig(name="tf_test", 
                             model=model, 
                             train_set=(x_train, y_train), 
                             test_set=(x_test, y_test),
                             hparams=model_parameters
                            )

5. Run distributed training

Finally, we are ready to launch the maggy experiment. You just need to pass 2 parameters: the training function and the configuration variable we defined in the previous steps.

experiment.lagom(training_function, config)