Learning Hierarchical Graph Neural Networks for Image Clustering

Abstract

We propose a hierarchical graph neural network (GNN) model that learns how to cluster a set of images into an unknown number of identities using a training set of images annotated with labels belonging to a disjoint set of identities. Our hierarchical GNN uses a novel approach to merge connected components predicted at each level of the hierarchy to form a new graph at the next level. Unlike fully unsupervised hierarchical clustering, the choice of grouping and complexity criteria stems naturally from supervision in the training set. The resulting method, Hi-LANDER, achieves an average of 54% improvement in F-score and 8% increase in Normalized Mutual Information (NMI) relative to current GNN-based clustering algorithms. Additionally, state-of-the-art GNN-based methods rely on separate models to predict linkage probabilities and node densities as intermediate steps of the clustering process. In contrast, our unified framework achieves a seven-fold decrease in computational cost. We release our training and inference code at https://github.com/dmlc/dgl/tree/master/examples/pytorch/hilander.

References in corpus (5)

• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• Semi-Supervised Classification with Graph Convolutional Networks
• Spatial Temporal Graph Convolutional Networks for Skeleton-Based Action Recognition
• FastGCN: Fast Learning with Graph Convolutional Networks via Importance Sampling
• GNN3DMOT: Graph Neural Network for 3D Multi-Object Tracking with Multi-Feature Learning