mlx_graphs.datasets.IMDB

class mlx_graphs.datasets.IMDB(base_dir: str | None = None, transform: Callable | None = None, pre_transform: Callable | None = None)

A subset of the Internet Movie Database (IMDB), as collected in the “MAGNN: Metapath Aggregated Graph Neural Network for Heterogeneous Graph Embedding” paper. IMDB is a heterogeneous graph containing three types of entities - movies (4,278 nodes), actors (5,257 nodes), and directors (2,081 nodes). The movies are divided into three classes (action, comedy, drama) according to their genre. Movie features correspond to elements of a bag-of-words representation of its plot keywords.

Parameters:

• base_dir (Optional[str]) – directory where the dataset should be saved.

• transform (Optional[Callable]) – A function/transform that takes in an HeteroGraphData object and returns a transformed version. The data object will be transformed before every access. (default: None)

• pre_transform (Optional[Callable]) – A function/transform that takes in an HeteroGraphData object and returns a transformed version. The data object will be transformed before being saved to disk. (default: None)

__init__(base_dir: str | None = None, transform: Callable | None = None, pre_transform: Callable | None = None)

Methods

__init__([base_dir, transform, pre_transform])
download()	Download the dataset at self.raw_path.
load()	Load the processed dataset
process()	Process the dataset and store data in self.data
save()	Save the processed dataset

Attributes

name	Name of the dataset
num_edge_classes	Returns a dictionary of the number of edge classes for each edge type.
num_edge_features	Returns a dictionary of the number of edge features for each edge type.
num_edges	Returns a dictionary of the number of edges for each edge type.
num_graph_features	Returns the number of graph features.
num_items	Returns the number of items in the dataset.
num_node_classes	Returns a dictionary of the number of node classes for each node type.
num_node_features	Returns a dictionary of the number of node features for each node type.
num_nodes	Returns a dictionary of the number of nodes for each node type.
processed_path	The path where processed files are stored.
raw_file_names
raw_path	The path where raw files are stored.