IMDB4k

class dhg.data.IMDB4k(data_root=None)[source]

Bases: dhg.data.base.BaseData

The IMDB-4k dataset is a movie dataset for node classification task. The dataset is an online database about movies and television programs, including information such as cast, production crew, and plot summaries. This is a subset of IMDB scraped from online, containing 4278 movies, 2081 directors, and 5257 actors after data preprocessing. Movies are labeled as one of three classes (Action, Comedy, and Drama) based on their genre information. Each movie is also described by a bag-of-words representation of its plot keywords. The vertice denotes author, and two types of correlation (co-director, co-actor) can be used for building hyperedges. More details see the MAGNN: Metapath Aggregated Graph Neural Network for Heterogeneous Graph Embedding paper.

The content of the IMDB-4k dataset includes the following:

  • num_classes: The number of classes: \(3\).

  • num_vertices: The number of vertices: \(4,278\).

  • num_director_edges: The number of hyperedges constructed by the co-director correlation: \(2,081\).

  • num_actor_edges: The number of hyperedges constructed by the co-actor correlation: \(5,257\).

  • dim_features: The dimension of movie features: \(3,066\).

  • features: The movie feature matrix. torch.Tensor with size \((4,278 \times 3,066)\).

  • labels: The label list. torch.LongTensor with size \((4,278, )\).

  • edge_by_director: The hyperedge list constructed by the co-director correlation. List with length \((2,081)\).

  • edge_by_actor: The hyperedge list constructed by the co-actor correlation. List with length \((5,257)\).

Parameters

data_root (str, optional) – The data_root has stored the data. If set to None, this function will auto-download from server and save into the default direction ~/.dhg/datasets/. Defaults to None.