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 on November 11, 2024.