Model Selection with Model Zoo via Graph Learning
Ziyu Li Hilco van der Wilk Danning Zhan Megha Khosla Alessandro Bozzon Rihan Hai
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.