Skip to content

Commit 277a747

Browse files
RobJYRobJY
committed
checkpoint: works up to step 3
1 parent 01c9020 commit 277a747

1 file changed

Lines changed: 238 additions & 0 deletions

File tree

‎docs/hpc/08_ml_ai_hpc/04_tensorflow.md‎

Lines changed: 238 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,238 @@
1+
# TensorFlow
2+
3+
:::info
4+
This was adapted from [Princeton University Multi-GPU Training with PyTorch](https://github.com/PrincetonUniversity/multi_gpu_training)
5+
:::
6+
7+
The starting point for [multi-GPU training with Keras](https://www.tensorflow.org/tutorials/distribute/keras) is `tf.distribute.MirroredStrategy`. In this approach, the model is copied to `N` GPUs and gradients are synced as we saw previously. Be sure to use [`tf.data`](https://www.tensorflow.org/api_docs/python/tf/data) to handle data loading as is done in the example on this page and is explained graphically [here](https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/guide/data_performance.ipynb#scrollTo=i3NtGI3r-jLp).
8+
9+
## Single-Node, Synchronous, Multi-GPU Training
10+
11+
Here were train the ResNet-50 model on the Cassava dataset (see [video](https://www.youtube.com/watch?v=xzSCvXDcX68) on TensorFlow YouTube channel). Here is another example [video](https://www.youtube.com/watch?v=HCLmM1PyDIs) using the "cats vs. dog" dataset.
12+
13+
### Step 1: Create a TensorFlow Overlay File
14+
15+
```bash
16+
# create a working directory
17+
[NetID@log-1 ~]$ mkdir /scratch/<NetID>/tensorflow-example
18+
[NetID@log-1 ~]$ cd /scratch/<NetID>/tensorflow-example
19+
20+
# copy over an overlay file with sufficient resources and unzip it
21+
[NetID@cm001 tensorflow-example]$ cp -rp /scratch/work/public/overlay-fs-ext3/overlay-15GB-500K.ext3.gz .
22+
[NetID@cm001 tensorflow-example]$ gunzip overlay-15GB-500K.ext3.gz
23+
24+
# start the singularity environment
25+
[NetID@cm001 tensorflow-example]$ singularity exec --overlay overlay-15GB-500K.ext3:rw /scratch/work/public/singularity/cuda12.1.1-cudnn8.9.0-devel-ubuntu22.04.2.sif /bin/bash
26+
27+
# install miniforge in singularity environment
28+
Singularity> wget --no-check-certificate https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-x86_64.sh
29+
Singularity> bash Miniforge3-Linux-x86_64.sh -b -p /ext3/miniforge3
30+
31+
# create an miniconda environment file
32+
Singularity> touch /ext3/env.sh
33+
Singularity> nano /ext3/env.sh
34+
# and add the following content to it:
35+
```
36+
37+
```bash
38+
#!/bin/bash
39+
40+
unset -f which
41+
42+
source /ext3/miniforge3/etc/profile.d/conda.sh
43+
export PATH=/ext3/miniforge3/bin:$PATH
44+
export PYTHONPATH=/ext3/miniforge3/bin:$PATH
45+
```
46+
47+
```bash
48+
# activate the new environment and initialize
49+
Singularity> source /ext3/env.sh
50+
Singularity> conda config --remove channels defaults
51+
Singularity> conda clean --all --yes
52+
Singularity> conda update -n base conda -y
53+
Singularity> conda install pip -y
54+
Singularity> conda install ipykernel -y
55+
56+
# check that you're using the miniconda environment
57+
# you should get the following output
58+
Singularity> which conda
59+
/ext3/miniforge3/bin/conda
60+
Singularity> which python
61+
/ext3/miniforge3/bin/python
62+
Singularity> which pip
63+
/ext3/miniforge3/bin/pip
64+
65+
# install TensorFlow
66+
Singularity> pip install tensorflow[and-cuda] tensorflow_datasets
67+
```
68+
69+
### Step 2: Download the Data
70+
71+
This example using the `cassava` dataset which requires 4 GB of storage space. Be sure to save this in your `/scratch` space and not in `/home`.
72+
73+
Please save the following into a file named `download_data_and_weights.py`:
74+
```python
75+
import tensorflow as tf
76+
import tensorflow_datasets as tfds
77+
78+
# download the data (4 GB) on the login node
79+
_ = tfds.load(name='cassava', with_info=True, as_supervised=True, data_dir='.')
80+
81+
# download the model weights on the login node
82+
_ = tf.keras.applications.ResNet50(weights="imagenet", include_top=False)
83+
```
84+
85+
Run the command below to download the data (4 GB in size):
86+
87+
```bash
88+
# switch to a data transfer node
89+
[NetID@log-1 tensorflow_example]$ ssh gdtn
90+
[NetID@dtn-1 ~]$ cd /scratch/NetID/tensorflow_example
91+
[NetID@dtn-1 tensorflow_example]$ singularity exec --nv --overlay /scratch/NetID/pytorch-example/my_pytorch.ext3:ro /scratch/work/public/singularity/cuda12.1.1-cudnn8.9.0-devel-ubuntu22.04.2.sif /bin/bash -c "source /ext3/env.sh; python download_data_and_weights.py"
92+
```
93+
94+
### Step 3: Inspect the Script
95+
96+
Below is the contents of `mnist_classify.py`:
97+
98+
```python
99+
import argparse
100+
import os
101+
import tensorflow_datasets as tfds
102+
import tensorflow as tf
103+
from time import perf_counter
104+
105+
def preprocess_data(image, label):
106+
image = tf.image.resize(image, (300, 300))
107+
image = tf.cast(image, tf.float32) / 255.0
108+
return image, label
109+
110+
def create_dataset(batch_size_per_replica, datasets, strategy):
111+
batch_size = batch_size_per_replica * strategy.num_replicas_in_sync
112+
return datasets['train'].map(preprocess_data, num_parallel_calls=tf.data.AUTOTUNE) \
113+
.cache() \
114+
.shuffle(1000) \
115+
.batch(batch_size) \
116+
.prefetch(tf.data.AUTOTUNE)
117+
118+
def create_model(num_classes):
119+
base_model = tf.keras.applications.ResNet50(weights="imagenet", include_top=False)
120+
x = base_model.output
121+
x = tf.keras.layers.GlobalAveragePooling2D()(x)
122+
x = tf.keras.layers.Dense(1016, activation="relu")(x)
123+
predictions = tf.keras.layers.Dense(num_classes, activation="softmax")(x)
124+
model = tf.keras.Model(inputs=base_model.input, outputs=predictions)
125+
return model
126+
127+
def train(epochs, num_classes, train_dataset, strategy):
128+
with strategy.scope():
129+
model = create_model(num_classes)
130+
model.compile(loss='sparse_categorical_crossentropy',
131+
optimizer=tf.keras.optimizers.Adam(learning_rate=0.0001),
132+
metrics=['accuracy'])
133+
134+
start_time = perf_counter()
135+
model.fit(train_dataset, epochs=epochs)
136+
print("Training time:", perf_counter() - start_time)
137+
return None
138+
139+
def print_info(num_replicas_in_sync, batch_size_per_replica, info, num_classes):
140+
print(f'TF Version: {tf.__version__}')
141+
print(f'Number of GPUs: {num_replicas_in_sync}')
142+
print(f'Batch size per GPU: {batch_size_per_replica}')
143+
print(f'Train records: {info.splits["train"].num_examples}')
144+
print(f'Test records: {info.splits["test"].num_examples}')
145+
print(f'Number of classes: {num_classes}')
146+
return None
147+
148+
if __name__ == '__main__':
149+
parser = argparse.ArgumentParser(description='Multi-GPU Training Example')
150+
parser.add_argument('--batch-size-per-replica', type=int, default=32, metavar='N',
151+
help='input batch size per GPU for training (default: 32)')
152+
parser.add_argument('--epochs', type=int, default=15, metavar='N',
153+
help='number of epochs to train (default: 15)')
154+
args = parser.parse_args()
155+
156+
datasets, info = tfds.load(name='cassava', with_info=True, as_supervised=True, data_dir=".")
157+
num_classes = info.features["label"].num_classes
158+
159+
strategy = tf.distribute.MirroredStrategy()
160+
train_dataset = create_dataset(args.batch_size_per_replica, datasets, strategy)
161+
train(args.epochs, num_classes, train_dataset, strategy)
162+
163+
print_info(strategy.num_replicas_in_sync, args.batch_size_per_replica, info, num_classes)
164+
```
165+
166+
### Step 4: Submit the Job
167+
168+
Below is a sample Slurm script:
169+
170+
```bash
171+
#!/bin/bash
172+
#SBATCH --job-name=cassava # create a short name for your job
173+
#SBATCH --nodes=1 # node count
174+
#SBATCH --ntasks=1 # total number of tasks across all nodes
175+
#SBATCH --cpus-per-task=16 # cpu-cores per task (>1 if multi-threaded tasks)
176+
#SBATCH --mem=64G # total memory per node (4G per cpu-core is default)
177+
#SBATCH --gres=gpu:2 # number of gpus per node
178+
#SBATCH --time=00:20:00 # total run time limit (HH:MM:SS)
179+
#SBATCH --mail-type=begin # send email when job begins
180+
#SBATCH --mail-type=end # send email when job ends
181+
#SBATCH --mail-user=<YourNetID>@princeton.edu
182+
183+
module purge
184+
module load anaconda3/2021.11
185+
conda activate tf2-v100
186+
187+
python cassava_classify.py --batch-size-per-replica=32 --epochs=15
188+
```
189+
190+
Note that `srun` is not called and there is only one task. Submit the job as follows:
191+
192+
```
193+
(tf2-v100) $ sbatch job.slurm
194+
```
195+
196+
### Performance
197+
198+
The training time is shown below for different choices of `cpus-per-task` and the number of GPUs:
199+
200+
| nodes | ntasks | cpus-per-task | GPUs | Training Time (s) | Mean GPU Utilization (%) |
201+
|:-------------:|:-------------:|:------------:|:--------:|:-----------------:|:-------------------------:|
202+
| 1 | 1 | 2 | 1 | 574 | 85 |
203+
| 1 | 1 | 4 | 1 | 565 | 83 |
204+
| 1 | 1 | 8 | 1 | 562 | 89 |
205+
| 1 | 1 | 16 | 1 | 564 | 90 |
206+
| 1 | 1 | 4 | 2 | 339 | 76 |
207+
| 1 | 1 | 8 | 2 | 334 | 81 |
208+
| 1 | 1 | 16 | 2 | 332 | 74 |
209+
| 1 | 1 | 4 | 3 | 256 | 68 |
210+
| 1 | 1 | 8 | 3 | 251 | 73 |
211+
| 1 | 1 | 16 | 3 | 249 | 66 |
212+
| 1 | 1 | 4 | 4 | 226 | 59 |
213+
| 1 | 1 | 8 | 4 | 220 | 58 |
214+
| 1 | 1 | 16 | 4 | 214 | 65 |
215+
| 1 | 1 | 32 | 4 | 218 | 63 |
216+
217+
"Training Time" in the table above is the time to run `model.fit(train_dataset, epochs=epochs)`. "Mean GPU utilization" was taken from the output of the `jobstats` command.
218+
219+
The figure below shows the speed-up as a function of the number of GPUs. The dashed line shows the maximum possible speed-up.
220+
221+
<img src="speedup_vs_gpus.png" alt="speed-up" width="700"/>
222+
223+
We see that linear scaling is not observed. That is, the training time when using 2 GPUs is not 1/2 of the training time when using one. To improve on this one would profile the script and identify the performance bottleneck. Some of the training images are 500 pixels wide. It could be that the preprocessing step is the slowest.
224+
225+
All runs were done on adroit-h11g1 while making certain that no other jobs were running on the node:
226+
227+
```
228+
#SBATCH --mem=770000M
229+
#SBATCH --nodelist=adroit-h11g1
230+
```
231+
232+
## Multi-node Training
233+
234+
Look to [MultiWorkerMirroredStrategy](https://www.tensorflow.org/guide/distributed_training#multiworkermirroredstrategy) for using the GPUs on more than one compute node. There is an example for the [Keras API](https://www.tensorflow.org/tutorials/distribute/multi_worker_with_keras). Consider using Horovod instead of this approach (see below).
235+
236+
## Horovod
237+
238+
[Horovod](https://horovod.ai/) is a distributed deep learning training framework for TensorFlow, Keras, PyTorch, and Apache MXNet. It is based on MPI.

0 commit comments

Comments
 (0)