Contact: hanqiu.chen@ece.gatech.edu
ECE, Georgia Institute of Technology
Rapid-INR: Storage Efficient CPU-free DNN Training Using Implicit Neural Representation
Hanqiu Chen, Hang Yang§, Stephen Fitzmeyer§ and Cong (Callie) Hao
(§Equal contribution)
Sharc Lab
Challenges in Computer Vision Task Training
2
Large data communication overhead and unsaturated computing resource
Training dataset is getting larger
(Hennessy, John L., and David A. Patterson. Computer architecture: a quantitative approach)
Storage Explosion
Memory Wall
Challenges in Computer Vision Task Training
3
Storage Explosion
Memory Wall
Implicit neural representation!
How to reduce training dataset storage and thus alleviate memory wall?
Implicit Neural Representation (INR)
4
RGB image
Neural Network for INR
Inputs to the INR NN: pixel coordinates
Outputs: rgb values
Implicit Neural Representation (INR)
5
RGB image
Neural Network for INR
Inputs to the INR NN: pixel coordinates
Outputs: rgb values
Why is INR beneficial for training?
INR Benefits Training Because…
6
Training dataset in JPEG format (e.g., ImageNet)
Training dataset in RGB format
Dedicated decoding
(CPU/GPU)
Many rounds of data movement to GPU
INR Benefits Training Because…
7
Training dataset in JPEG format (e.g., ImageNet)
Training dataset in RGB format
Dedicated decoding
(CPU/GPU)
Many rounds of data movement to GPU
Single round to GPU
Easy Decoding!
Training dataset in INR format (i.e., NN model weights)
ImageNet can be compressed from 138GB to 14GB
138GB
~2TB
14GB
INR Benefits Training Because…
8
Training dataset in JPEG format (e.g., ImageNet)
Training dataset in RGB format
Dedicated decoding
(CPU/GPU)
Many rounds of data movement to GPU
Single round to GPU
Easy Decoding!
Training dataset in INR format (i.e., NN model weights)
138GB
~2TB
14GB
Compact and general
Storage efficient
Continuous and smooth
Previous Works and Motivations
9
INR models | Advantages | Limitations |
COIN[1] | The pioneer work using INR for image encoding | Cannot beat JPEG if high reconstruction quality is needed |
COIN++[2] | Use meta-learning to improve reconstruction quality | 3x slower decoding speed compared with COIN |
[1] Dupont, Emilien, Adam Goliński, Milad Alizadeh, Yee Whye Teh, and Arnaud Doucet. "Coin: Compression with implicit neural representations." arXiv:2103.03123 (2021).
[2] Dupont, Emilien, Hrushikesh Loya, Milad Alizadeh, Adam Goliński, Yee Whye Teh, and Arnaud Doucet. "COIN++: Neural compression across modalities." arXiv:2201.12904 (2022).
Previous Works and Motivations
10
INR models | Advantages | Limitations |
COIN[1] | The pioneer work using INR for image encoding | Cannot beat JPEG if high reconstruction quality is needed |
COIN++[2] | Use meta-learning to improve reconstruction quality | 3x slower decoding speed compared with COIN |
[1] Dupont, Emilien, Adam Goliński, Milad Alizadeh, Yee Whye Teh, and Arnaud Doucet. "Coin: Compression with implicit neural representations." arXiv:2103.03123 (2021).
[2] Dupont, Emilien, Hrushikesh Loya, Milad Alizadeh, Adam Goliński, Yee Whye Teh, and Arnaud Doucet. "COIN++: Neural compression across modalities." arXiv:2201.12904 (2022).
Our Rapid-INR
Can beat JPEG in terms of both storage efficiency and reconstruction quality
Can achieve high decoding speed which is naturally suitable for GPUs
Rapid-INR v.s. Other Pipelines
11
�
�
NVIDIA DALI: https://docs.nvidia.com/deeplearning/dali/user-guide/docs/pipeline.html
PyTorch
RGB
Disk
Many batches repeatedly
One batch
DALI
Disk
Decoding… (slow)
Decoding… (faster)
RGB
Rapid-INR
Disk
Decoding… (fastest!)
RGB
INR
Copy once
No CPU!
Contributions
12
�
�
CPU-free training with exceptional speedup.
Rapid-INR
High decoding parallelism without specialized hardware
Optimized compression method for efficient storage.
Ease of use with high generality.
Up to 6x speedup compared with PyTorch training framework baseline.
Achieve pixel-level parallelism on general GPUs.
Dynamic pruning and layer-wise quantization for INR weights helps reduce 95% storage space.
Spatial resolution agnostic decoding allows for flexible and easy integration with mainstream CV training pipelines.
INR image compression
Training acceleration
Rapid-INR Key Techniques
13
1
Offline INR encoder for image compression: dataset only needed to be encoded once, then can be applied to multiple tasks.
2
Highly parallelized and low latency on-device INR decoder without specialized hardware.
3
Dynamic pruning based on reconstructed image quality.
Layer-wise quantization based on INR weights distribution of different layers. We quantize the hidden layers to 8 bits.
4
5
Design space exploration for optimized INR architectures.
Rapid-INR Key Techniques
14
1
Offline INR encoder for image compression: dataset only needed to be encoded once, then can be applied to multiple tasks.
2
Highly parallelized and low latency on-device INR decoder without specialized hardware.
3
Dynamic pruning based on reconstructed image quality.
Layer-wise quantization based on INR weights distribution of different layers. We quantize the hidden layers to 8 bits.
4
5
Design space exploration for optimized INR architectures.
Rapid-INR Encoding Architecture
15
RGB images
Disk
Weights encoded
Each image encoded with different MLPs
Offline encoding: Datasets only need to be encoded once, then can be reusable for different tasks.
Flexible INR size: The size of INR changes with the input image size to achieve a balance between storage efficiency and reconstruction quality.
Rapid-INR Decoding Architecture
16
Disk
Copy once!
CUDA memory
RGB image buffer
INR decoding
Data transferring before training starts
Pixel-level parallelism
Accelerated augmentation
Images are decoded batch by batch during backbone training
Backbone training
Specialized hardware General GPUs
Rapid-INR Key Techniques
17
1
Offline INR encoder for image compression: dataset only needed to be encoded once, then can be applied to multiple tasks.
2
Highly parallelized and low latency on-device INR decoder without specialized hardware.
3
Dynamic pruning based on reconstructed image quality.
Layer-wise quantization based on INR weights distribution of different layers. We quantize the hidden layers to 8 bits.
4
5
Design space exploration for optimized INR architectures.
Dynamic Pruning
18
Dynamic pruning: Different pruning ratios for images with different reconstruction quality (PSNR).
Different color richness and variety
Low PSNR images
Different PSNR
Reduce backbone training accuracy
Larger INR
Higher PSNR
Three empirical observations
Pruning Ratio Selection
19
Close to linear relationship
Findings
When pruning ratio increases from 0.1 ~ 0.8, the PSNR decreases linearly.
Pruning and Quantization Evaluation
20
�
�
Pruning and Quantization Evaluation
21
�
�
JPEG quality and size
Pruning and Quantization Evaluation
22
�
�
JPEG quality and size
Rapid-INR quality and size with quantization and pruning
Pruning and Quantization Evaluation
23
�
�
JPEG quality and size
Rapid-INR quality and size with quantization and pruning
Rapid-INR Key Techniques
24
1
Offline INR encoder for image compression: dataset only needed to be encoded once, then can be applied to multiple tasks.
2
Highly parallelized and low latency on-device INR decoder without specialized hardware.
3
Dynamic pruning based on reconstructed image quality.
Layer-wise quantization based on INR weights distribution of different layers. We quantize the hidden layers to 8 bits.
4
5
Design space exploration for optimized INR architectures.
Neural Architecture Search
25
We use less than 10 layers MLPs with uniformed structure (same hidden layers dimension)
……
Hidden layers with same dimension
uniformed structure
Uniformed structure MLP
Image Reconstruction Quality
26
highest quality
Backbone Training Accuracy
27
14.1%
0.3%
2.5%
6.2%
3.4%
4%
Backbone Training Speedup
28
Backbone training time comparison across different pipelines and datasets using different numbers of CPU threads for preprocessing.
1.95x
6x
4.4x
1.2x
Summary & Thanks
29
Contact: hchen799@gatech.edu Sharc-lab @ Georgia Tech (https://sharclab.ece.gatech.e du/)
Thank Xuebin Yao and Pradeep Subedi from Samsung for their insightful discussions!