https://github.com/sseung0703/tf2-jit-compile-on-multi-gpu
Tensorflow2 training code with jit compiling on multi-GPU.
https://github.com/sseung0703/tf2-jit-compile-on-multi-gpu
jit-compilation multi-gpu-training tensorflow2 xla
Last synced: 5 months ago
JSON representation
Tensorflow2 training code with jit compiling on multi-GPU.
- Host: GitHub
- URL: https://github.com/sseung0703/tf2-jit-compile-on-multi-gpu
- Owner: sseung0703
- Created: 2020-12-23T13:15:44.000Z (almost 5 years ago)
- Default Branch: main
- Last Pushed: 2021-01-28T07:44:03.000Z (over 4 years ago)
- Last Synced: 2025-03-31T01:51:13.302Z (6 months ago)
- Topics: jit-compilation, multi-gpu-training, tensorflow2, xla
- Language: Python
- Homepage:
- Size: 157 KB
- Stars: 17
- Watchers: 1
- Forks: 2
- Open Issues: 0
-
Metadata Files:
- Readme: README.md
Awesome Lists containing this project
README
## Tensorflow2 with jit (xla) compiling on multi-gpu training.
- CIFAR and ILSVRC training code with **jit compiling** and **distributed learning** on the multi-GPU system.
- I highly recommend using Jit compiling because most of the algorithm is static and can be compiled, which gives memory usage reduction and training speed improvement.
- This repository is built by **custom layers** and **custom training loop** for my project, but if you only want to check how to use jit compiling with distributed learning, check 'train.py' and 'op_util.py'.
## Requirement
- **Tensorflow >= 2.5**
- Pillow
## How to run
- ILSVRC
```
python train.py --compile --gpu_id {} --dataset ILSVRC --data_path /path/to/your/ILSVRC/home --train_path /path/to/log
```
- CIFAR{10,100}
```
python train.py --compile --gpu_id {} --dataset CIFAR{10,100} --train_path /path/to/log
```
## Experimental results
- I used four 1080ti.
- Jit compiling gives a 40% speedup for training time.
| | Accuracy | Training time
|------------| ------------- | -------------
|Distributed only | 75.83 | 94.61
|Distributed with Jit | 75.57 | 56.98
