A Structure-Aware Framework for Learning Device Placements on Computation Graphs

Shukai Duan^1*, Heng Ping^1*, Nikos Kanakaris^1*, Xiongye Xiao^1*, Panagiotis Kyriakis^2*, Nesreen K. Ahmed³, Peiyu Zhang¹, Guixiang Ma⁴, Mihai Capota⁴, Shahin Nazarian¹, Theodore L. Willke⁴, Paul Bogdan¹

¹University of Southern California, ²Meta, ³Cisco Outshift, ⁴Intel Labs

---

This repository is the official implementation of the paper "A Structure-Aware Framework for Learning Device Placements on Computation Graphs" accepted at the NeurIPS 2024 conference.

Overview

• Proposes a novel framework for device placement, using smaller computation graphs extracted from the OpenVINO toolkit and reinforcement learning.
• Consists of five steps, including graph coarsening, node representation learning, and policy optimization.
• Supports end-to-end training and considers the directed and acyclic nature of the computation graphs.
• A model variant inspired by graph parsing networks and complex network analysis, enabling joint graph representation learning and personalized graph partitioning.

Setup

• Code is written in Python 3.9
• Set up the conda environment with following command: conda env create -f environment.yml

Instructions

To use the RNN baseline:

1. Open RNN_bert.ipynb and RNN_resnet_inceptionv3.ipynb in a jupyter notebook.
2. Edit the input data shape specified by your model to match your XML file.
3. Run all cells in each notebook.

To use the Placeto baseline:

1. Unzip PlacetoOpenVINO.zip under your directory and update the location of your XML file and the appropriate input.
2. To run the code, install and use our TensorFlow and OpenVINO Conda environment by executing conda env create -f environment_tf.yml.

To run our approach, HSDAG:

1. Provide the location of your XML file in config.py.
2. Run the main.py script with the appropriate input parameters for your experiment.

Notes:

1. Our source code for Placeto is modified based on https://github.com/mmitropolitsky/device-placement.
2. In HSDAG, we use the Graph Parsing Network. Visit the repository to check out their license.

Blogs

• Meta
• USC Viterbi school
• Intel labs
• Cisco outshift
• Papers with code

Bibliography

Please cite our paper if you use this code:

@misc{duan2024structureawareframeworklearningdevice,
      title={A structure-aware framework for learning device placements on computation graphs}, 
      author={Shukai Duan and Heng Ping and Nikos Kanakaris and Xiongye Xiao and Peiyu Zhang and Panagiotis Kyriakis and Nesreen K. Ahmed and Guixiang Ma and Mihai Capota and Shahin Nazarian and Theodore L. Willke and Paul Bogdan},
      year={2024},
      eprint={2405.14185},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2405.14185}, 
}