torch_geometric.datasets.StochasticBlockModelDataset
- class StochasticBlockModelDataset(root: str, block_sizes: Union[List[int], Tensor], edge_probs: Union[List[List[float]], Tensor], num_channels: Optional[int] = None, is_undirected: bool = True, transform: Optional[Callable] = None, pre_transform: Optional[Callable] = None, **kwargs)[source]
Bases:
InMemoryDataset
A synthetic graph dataset generated by the stochastic block model. The node features of each block are sampled from normal distributions where the centers of clusters are vertices of a hypercube, as computed by the
sklearn.datasets.make_classification()
method.- Parameters
root (str) – Root directory where the dataset should be saved.
block_sizes ([int] or LongTensor) – The sizes of blocks.
edge_probs ([[float]] or FloatTensor) – The density of edges going from each block to each other block. Must be symmetric if the graph is undirected.
num_channels (int, optional) – The number of node features. If given as
None
, node features are not generated. (default:None
)is_undirected (bool, optional) – Whether the graph to generate is undirected. (default:
True
)transform (callable, optional) – A function/transform that takes in an
torch_geometric.data.Data
object and returns a transformed version. The data object will be transformed before every access. (default:None
)pre_transform (callable, optional) – A function/transform that takes in an
torch_geometric.data.Data
object and returns a transformed version. The data object will be transformed before being saved to disk. (default:None
)**kwargs (optional) – The keyword arguments that are passed down to the
sklearn.datasets.make_classification()
method for drawing node features.