Transferability Prediction for Model Recommendation: A Graph Learning Method

The increasing availability of pre-trained deep learning models has made transfer learning a popular approach for improving the performance of target tasks. However, selecting the most suitable pre-trained model from a large pool of candidates, especially with heterogeneous architectures, remains a significant challenge. In this paper, we address the problem of model recommendation by predicting the transferability of pre-trained models to a target dataset. Specifically, we formulate this transferability prediction task as a bipartite graph learning problem, where models and datasets are represented as nodes, and the edges represent the transferability between them. This formulation allows us to leverage model and dataset metadata as node features and empirical transferability as edge weights. A customized graph attention network, termed ModelRecommend, is then employed to estimate the transferability between models and datasets, even without access to source data. By learning low-dimensional model and dataset embeddings from their metadata and predicting edges accordingly in the latent space, our method effectively captures intrinsic model and dataset interrelationships, enabling transferability prediction for unseen model-dataset pairs. We experimentally demonstrate the effectiveness of ModelRecommend in both general image classification and segmentation tasks in autonomous driving scenarios, highlighting improvements of 15% and 12% compared to state-of-the-art methods, respectively. The code is available in https://github.com/xychendave/ModelRecommendation.