This is a simple demo of Tunable Subnetwork Splitting Method (TSSM) as described in our paper:
Junxiang Wang, Zheng Chai, Yue Cheng, and Liang Zhao. Tunable Subnetwork Splitting for Model-parallelism of Neural Network Training. (ICML 2020 Workshop: Beyond First Order Methods in Machine Learning)
python setup.py install
tensorflow 2.1.0
numpy 1.18.1
scipy 1.4.1
keras 2.3.1
tensorflow_datasets 1.3.2
python main.py
Three benchmark datasets MNIST, Fashion-MNIST, kMNIST are included in this package.
Please cite our paper if you use this code in your own work:
@article{wang2020,
author = {Wang, Junxiang and Zheng Chai and Yue Chen and Zhao, Liang},
title = {Tunable Subnetwork Splitting for Model-parallelism of Neural Network Training},
year = {2020},
booktitle = {ICML 2020 Workshop: Beyond First Order Methods in Machine Learning},
}