Transfer Learning Across Heterogeneous Features For Efficient Tensor Program Generation
This work addresses the efficiency of tensor program optimization for machine learning practitioners, offering incremental improvements in tuning speed and resource usage.
The paper tackles the challenge of tuning tensor program generation for heterogeneous hardware by learning and transferring joint neural network and hardware features, achieving up to 45% dataset pruning without loss in accuracy and reducing tuning time by 25-40% while maintaining or improving inference performance.
Tuning tensor program generation involves searching for various possible program transformation combinations for a given program on target hardware to optimize the tensor program execution. It is already a complex process because of the massive search space and exponential combinations of transformations make auto-tuning tensor program generation more challenging, especially when we have a heterogeneous target. In this research, we attempt to address these problems by learning the joint neural network and hardware features and transferring them to the new target hardware. We extensively study the existing state-of-the-art dataset, TenSet, perform comparative analysis on the test split strategies and propose methodologies to prune the dataset. We adopt an attention-inspired approach for tuning the tensor programs enabling them to embed neural network and hardware-specific features. Our approach could prune the dataset up to 45\% of the baseline without compromising the Pairwise Comparison Accuracy (PCA). Further, the proposed methodology can achieve on-par or improved mean inference time with 25%-40% of the baseline tuning time across different networks and target hardware.