A Structure-Aware Framework for Learning Device Placements on Computation Graphs
Published in NeurIPS 2024, 2025
We propose a structure-aware framework for optimizing device placement on computation graphs used in neural network inference. The method integrates graph coarsening, node representation learning, and reinforcement learning–based policy optimization into an end-to-end trainable pipeline. Experiments on Inception-V3, ResNet, and BERT demonstrate substantial speedups—up to 58.2% over CPU execution and up to 60.24% compared to existing baselines. Official implementation is available in the HSDAG GitHub repository.