title multigpu distributed training with tensorflow fchollethttps twitter comfchollet date created 20200428 last modified 20230629 description guide to multigpu training for keras models with tensorflow accelerator gpu introduction there are generally two ways to distribute computation across multiple devices data parallelism where a single model gets replicated on multiple devices or multiple machines each of them processes different batches of data then they merge their results there exist many variants of this setup that differ in how the different model replicas merge results in whether they stay in sync at every batch or whether they are more loosely coupled etc model parallelism where different parts of a single model run on different devices processing a single batch of data together this works best with models that have a naturallyparallel architecture such as models that feature multiple branches this guide focuses on data parallelism in particular synchronous data parallelism where the different replicas of the model stay in sync after each batch they process synchronicity keeps the model convergence behavior identical to what you would see for singledevice training specifically this guide teaches you how to use the tf distribute api to train keras models on multiple gpus with minimal changes to your code on multiple gpus typically 2 to 16 installed on a single machine single host multidevice training this is the most common setup for researchers and smallscale industry workflows setup singlehost multidevice synchronous training in this setup you have one machine with several gpus on it typically 2 to 16 each device will run a copy of your model called a replica for simplicity in what follows we ll assume we re dealing with 8 gpus at no loss of generality how it works at each step of training the current batch of data called global batch is split into 8 different subbatches called local batches for instance if the global batch has 512 samples each of the 8 local batches will have 64 samples each of the 8 replicas independently processes a local batch they run a forward pass then a backward pass outputting the gradient of the weights with respect to the loss of the model on the local batch the weight updates originating from local gradients are efficiently merged across the 8 replicas because this is done at the end of every step the replicas always stay in sync in practice the process of synchronously updating the weights of the model replicas is handled at the level of each individual weight variable this is done through a mirrored variable object how to use it to do singlehost multidevice synchronous training with a keras model you would use the tf distribute mirroredstrategy api https www tensorflow orgapidocspythontfdistributemirroredstrategy here s how it works instantiate a mirroredstrategy optionally configuring which specific devices you want to use by default the strategy will use all gpus available use the strategy object to open a scope and within this scope create all the keras objects you need that contain variables typically that means creating compiling the model inside the distribution scope in some cases the first call to fit may also create variables so it s a good idea to put your fit call in the scope as well train the model via fit as usual importantly we recommend that you use tf data dataset objects to load data in a multidevice or distributed workflow schematically it looks like this python create a mirroredstrategy strategy tf distribute mirroredstrategy print number of devices formatstrategy numreplicasinsync open a strategy scope with strategy scope everything that creates variables should be under the strategy scope in general this is only model construction compile model model model compile train the model on all available devices model fittraindataset validationdatavaldataset test the model on all available devices model evaluatetestdataset here s a simple endtoend runnable example make a simple 2layer denselyconnected neural network return the mnist dataset in the form of a tf data dataset preprocess the data these are numpy arrays reserve numvalsamples samples for validation create a mirroredstrategy open a strategy scope everything that creates variables should be under the strategy scope in general this is only model construction compile train the model on all available devices test the model on all available devices using callbacks to ensure fault tolerance when using distributed training you should always make sure you have a strategy to recover from failure fault tolerance the simplest way to handle this is to pass modelcheckpoint callback to fit to save your model at regular intervals e g every 100 batches or every epoch you can then restart training from your saved model here s a simple example prepare a directory to store all the checkpoints either restore the latest model or create a fresh one if there is no checkpoint available create a mirroredstrategy open a strategy scope and createrestore the model this callback saves a savedmodel every epoch we include the current epoch in the folder name running the first time creates the model calling the same function again will resume from where we left off tf data performance tips when doing distributed training the efficiency with which you load data can often become critical here are a few tips to make sure your tf data pipelines run as fast as possible note about dataset batching when creating your dataset make sure it is batched with the global batch size for instance if each of your 8 gpus is capable of running a batch of 64 samples you call use a global batch size of 512 calling dataset cache if you call cache on a dataset its data will be cached after running through the first iteration over the data every subsequent iteration will use the cached data the cache can be in memory default or to a local file you specify this can improve performance when your data is not expected to change from iteration to iteration you are reading data from a remote distributed filesystem you are reading data from local disk but your data would fit in memory and your workflow is significantly iobound e g reading decoding image files calling dataset prefetchbuffersize you should almost always call prefetchbuffersize after creating a dataset it means your data pipeline will run asynchronously from your model with new samples being preprocessed and stored in a buffer while the current batch samples are used to train the model the next batch will be prefetched in gpu memory by the time the current batch is over that s it introduction there are generally two ways to distribute computation across multiple devices data parallelism where a single model gets replicated on multiple devices or multiple machines each of them processes different batches of data then they merge their results there exist many variants of this setup that differ in how the different model replicas merge results in whether they stay in sync at every batch or whether they are more loosely coupled etc model parallelism where different parts of a single model run on different devices processing a single batch of data together this works best with models that have a naturally parallel architecture such as models that feature multiple branches this guide focuses on data parallelism in particular synchronous data parallelism where the different replicas of the model stay in sync after each batch they process synchronicity keeps the model convergence behavior identical to what you would see for single device training specifically this guide teaches you how to use the tf distribute api to train keras models on multiple gpus with minimal changes to your code on multiple gpus typically 2 to 16 installed on a single machine single host multi device training this is the most common setup for researchers and small scale industry workflows setup single host multi device synchronous training in this setup you have one machine with several gpus on it typically 2 to 16 each device will run a copy of your model called a replica for simplicity in what follows we ll assume we re dealing with 8 gpus at no loss of generality how it works at each step of training the current batch of data called global batch is split into 8 different sub batches called local batches for instance if the global batch has 512 samples each of the 8 local batches will have 64 samples each of the 8 replicas independently processes a local batch they run a forward pass then a backward pass outputting the gradient of the weights with respect to the loss of the model on the local batch the weight updates originating from local gradients are efficiently merged across the 8 replicas because this is done at the end of every step the replicas always stay in sync in practice the process of synchronously updating the weights of the model replicas is handled at the level of each individual weight variable this is done through a mirrored variable object how to use it to do single host multi device synchronous training with a keras model you would use the tf distribute mirroredstrategy api https www tensorflow org api_docs python tf distribute mirroredstrategy here s how it works instantiate a mirroredstrategy optionally configuring which specific devices you want to use by default the strategy will use all gpus available use the strategy object to open a scope and within this scope create all the keras objects you need that contain variables typically that means creating compiling the model inside the distribution scope in some cases the first call to fit may also create variables so it s a good idea to put your fit call in the scope as well train the model via fit as usual importantly we recommend that you use tf data dataset objects to load data in a multi device or distributed workflow schematically it looks like this python create a mirroredstrategy strategy tf distribute mirroredstrategy print number of devices format strategy num_replicas_in_sync open a strategy scope with strategy scope everything that creates variables should be under the strategy scope in general this is only model construction compile model model model compile train the model on all available devices model fit train_dataset validation_data val_dataset test the model on all available devices model evaluate test_dataset here s a simple end to end runnable example make a simple 2 layer densely connected neural network return the mnist dataset in the form of a tf data dataset preprocess the data these are numpy arrays reserve num_val_samples samples for validation create a mirroredstrategy open a strategy scope everything that creates variables should be under the strategy scope in general this is only model construction compile train the model on all available devices test the model on all available devices using callbacks to ensure fault tolerance when using distributed training you should always make sure you have a strategy to recover from failure fault tolerance the simplest way to handle this is to pass modelcheckpoint callback to fit to save your model at regular intervals e g every 100 batches or every epoch you can then restart training from your saved model here s a simple example prepare a directory to store all the checkpoints either restore the latest model or create a fresh one if there is no checkpoint available create a mirroredstrategy open a strategy scope and create restore the model this callback saves a savedmodel every epoch we include the current epoch in the folder name running the first time creates the model calling the same function again will resume from where we left off tf data performance tips when doing distributed training the efficiency with which you load data can often become critical here are a few tips to make sure your tf data pipelines run as fast as possible note about dataset batching when creating your dataset make sure it is batched with the global batch size for instance if each of your 8 gpus is capable of running a batch of 64 samples you call use a global batch size of 512 calling dataset cache if you call cache on a dataset its data will be cached after running through the first iteration over the data every subsequent iteration will use the cached data the cache can be in memory default or to a local file you specify this can improve performance when your data is not expected to change from iteration to iteration you are reading data from a remote distributed filesystem you are reading data from local disk but your data would fit in memory and your workflow is significantly io bound e g reading decoding image files calling dataset prefetch buffer_size you should almost always call prefetch buffer_size after creating a dataset it means your data pipeline will run asynchronously from your model with new samples being preprocessed and stored in a buffer while the current batch samples are used to train the model the next batch will be prefetched in gpu memory by the time the current batch is over that s it
import os

os.environ["KERAS_BACKEND"] = "tensorflow"

import tensorflow as tf
import keras




def get_compiled_model():
    
    inputs = keras.Input(shape=(784,))
    x = keras.layers.Dense(256, activation="relu")(inputs)
    x = keras.layers.Dense(256, activation="relu")(x)
    outputs = keras.layers.Dense(10)(x)
    model = keras.Model(inputs, outputs)
    model.compile(
        optimizer=keras.optimizers.Adam(),
        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        metrics=[keras.metrics.SparseCategoricalAccuracy()],
    )
    return model


def get_dataset():
    batch_size = 32
    num_val_samples = 10000

    
    (x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

    
    x_train = x_train.reshape(-1, 784).astype("float32") / 255
    x_test = x_test.reshape(-1, 784).astype("float32") / 255
    y_train = y_train.astype("float32")
    y_test = y_test.astype("float32")

    
    x_val = x_train[-num_val_samples:]
    y_val = y_train[-num_val_samples:]
    x_train = x_train[:-num_val_samples]
    y_train = y_train[:-num_val_samples]
    return (
        tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(
            batch_size
        ),
        tf.data.Dataset.from_tensor_slices((x_val, y_val)).batch(batch_size),
        tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(batch_size),
    )



strategy = tf.distribute.MirroredStrategy()
print("Number of devices: {}".format(strategy.num_replicas_in_sync))


with strategy.scope():
    
    
    model = get_compiled_model()

    
    train_dataset, val_dataset, test_dataset = get_dataset()
    model.fit(train_dataset, epochs=2, validation_data=val_dataset)

    
    model.evaluate(test_dataset)




checkpoint_dir = "./ckpt"
if not os.path.exists(checkpoint_dir):
    os.makedirs(checkpoint_dir)


def make_or_restore_model():
    
    
    checkpoints = [
        checkpoint_dir + "/" + name for name in os.listdir(checkpoint_dir)
    ]
    if checkpoints:
        latest_checkpoint = max(checkpoints, key=os.path.getctime)
        print("Restoring from", latest_checkpoint)
        return keras.models.load_model(latest_checkpoint)
    print("Creating a new model")
    return get_compiled_model()


def run_training(epochs=1):
    
    strategy = tf.distribute.MirroredStrategy()

    
    with strategy.scope():
        model = make_or_restore_model()

        callbacks = [
            
            
            keras.callbacks.ModelCheckpoint(
                filepath=checkpoint_dir + "/ckpt-{epoch}.keras",
                save_freq="epoch",
            )
        ]
        model.fit(
            train_dataset,
            epochs=epochs,
            callbacks=callbacks,
            validation_data=val_dataset,
            verbose=2,
        )



run_training(epochs=1)


run_training(epochs=1)