Multi GPU training with Tensorflow

Contents

Multi GPU training with Tensorflow#

Training on multiple GPUs is currently on the roadmap of keras 3.0 for both Tensorflow and Pytorch backend.
Until this is implemented in Keras 3.0, we don’t include multi GPU training inside the Crested trainer class, but you can still train on multiple GPU’s using the standard Tensorflow’s tf.distribute.MirroredStrategy.
You only need to wrap your model and optimizer creation and training inside the strategy.scope() context manager.
Data preparation is the same as in the single GPU case.

Example#

import tensorflow as tf

strategy = tf.distribute.MirroredStrategy()

datamodule = crested.tl.data.AnnDataModule(
        my_adata,
        genome=my_genome,
        chromsizes_file=my_chromsizes_file,
        batch_size=128,
        max_stochastic_shift=5,
        shuffle=True,
    )

with strategy.scope():
    model_architecture = crested.tl.zoo.chrombpnet(seq_len=2114, num_classes=4)

    config = crested.tl.default_configs("peak_regression")
    
    trainer = crested.tl.Crested(
            data=datamodule,
            model=model_architecture,
            config=config,
            project_name="test_project",
            run_name="test_run",
            logger="wandb",
        )

    trainer.fit(epochs=50)
Contents