Pre-trained deep learning (DL) models are increasingly accessible in public repositories, i.e., model zoos. Given a new prediction task, finding the best model to fine-tune can be computationally intensive and costly, especially when the number of pre-trained models is large. Selecting the right pre-trained models is crucial, yet complicated by the diversity of models from various model families (like ResNet, Vit, Swin) and the hidden relationships between models and datasets. Existing methods, which utilize basic information from models and datasets to compute scores indicating model performance on target datasets, overlook the intrinsic relationships, limiting their effectiveness in model selection. In this study, we introduce TransferGraph, a novel framework that reformulates model selection as a graph learning problem. TransferGraph constructs a graph using extensive metadata extracted from models and datasets, while capturing their inherent relationships. Through comprehensive experiments across 16 real datasets, both images and texts, we demonstrate TransferGraph's effectiveness in capturing essential model-dataset relationships, yielding up to a 32% improvement in correlation between predicted performance and the actual fine-tuning results compared to the state-of-the-art methods.

翻译：预训练深度学习（DL）模型在公共代码库（即模型动物园）中日益普及。针对新预测任务，从众多预训练模型中挑选最优模型进行微调，计算成本高且代价昂贵。选择合适的预训练模型至关重要，但受限于不同模型家族（如ResNet、ViT、Swin）的多样性以及模型与数据集间的隐含关系。现有方法利用模型和数据集的表层信息计算评分以评估模型在目标数据集上的表现，忽视了内在关联，限制了其在模型选择中的有效性。本研究提出TransferGraph这一新型框架，将模型选择重新定义为图学习问题。该框架通过提取模型和数据集的元数据构建图结构，并捕获其内在关系。在涵盖图像与文本的16个真实数据集上的综合实验表明，TransferGraph能有效捕捉模型与数据集的关键关联，与现有最优方法相比，预测性能与实际微调结果的相关性提升高达32%。