How to Use Keras-Based Multi-GPU
Keras is an open-source neural network library written in Python. It is a high-level neural network API that can run on top of MXNet, Deeplearning4j, TensorFlow, Microsoft Cognitive Toolkit, or Theano. In the NEURON system, queues like cas_v100_2, cas_v100nv_4, cas_v100nv_8, and amd_a100nv_8 are equipped with 2, 4, or 8 GPUs per node, providing an environment where multiple GPUs can be used for neural network training even within a single node.
A. Code modifications and job submission methods for using Multi-GPU
1. Add the [from keras.utils import multi_gpu_model] module
import keras
from keras.datasets import cifar10
from keras.models import Sequential
from keras.layers import Dense, Dropout, Activation, Flatten
from keras.layers import Conv2D, MaxPooling2D
from keras.utils import multi_gpu_model
import os
2. Declare the use of multi-GPU in the code
# initiate RMSprop optimizer
opt = keras.optimizers.rmsprop(lr=0.0001, decay=1e-6)
# multi-gpu
model = multi_gpu_model(model, gpus=2)
# Let's train the model using RMSprop
model.compile(loss='categorical_crossentropy', optimizer=opt, metrics=['accuracy'])
※ Set the number of GPUs to the desired amount. (For example, in a cas_v100nv_4 node, set gpus=4)
3. Job submission script
#!/bin/sh
#SBATCH -J keras
#SBATCH --time=24:00:00
#SBATCH -o %x_%j.out
#SBATCH -e %x_%j.err
#SBATCH -p ivy_v100_2
#SBATCH --comment tensorflow
#SBATCH --gres=gpu:2
#SBATCH -N 1
module purge
module load gcc/8.3.0 cuda/10.0 cudampi/openmpi-3.1.0 conda/tensorflow_1.13
srun python example.py
Last updated