wehub-resource-sync
54 lines
1.6 KiB
Python
54 lines
1.6 KiB
Python
import unittest
|
|
|
|
import backend as F
|
|
|
|
import dgl
|
|
from dgl.data.movielens import MovieLensDataset
|
|
|
|
|
|
@unittest.skipIf(
|
|
F._default_context_str == "gpu",
|
|
reason="Datasets don't need to be tested on GPU.",
|
|
)
|
|
@unittest.skipIf(
|
|
dgl.backend.backend_name != "pytorch", reason="only supports pytorch"
|
|
)
|
|
def test_movielens():
|
|
transform = dgl.AddSelfLoop(new_etypes=True)
|
|
|
|
movielens = MovieLensDataset(name="ml-100k", valid_ratio=0.2, verbose=True)
|
|
g = movielens[0]
|
|
assert g.num_edges("user-movie") == g.num_edges("movie-user") == 100000
|
|
assert (
|
|
g.nodes["user"].data["feat"].shape[1]
|
|
== g.nodes["user"].data["feat"].shape[1]
|
|
== g.nodes["user"].data["feat"].shape[1]
|
|
== 23
|
|
)
|
|
assert (
|
|
g.nodes["movie"].data["feat"].shape[1]
|
|
== g.nodes["movie"].data["feat"].shape[1]
|
|
== g.nodes["movie"].data["feat"].shape[1]
|
|
== 320
|
|
)
|
|
|
|
movielens = MovieLensDataset(
|
|
name="ml-100k", valid_ratio=0.2, transform=transform, verbose=True
|
|
)
|
|
g1 = movielens[0]
|
|
assert g1.num_edges() - g.num_edges() == g.num_nodes()
|
|
assert g1.num_edges() - g.num_edges() == g.num_nodes()
|
|
assert g1.num_edges() - g.num_edges() == g.num_nodes()
|
|
|
|
movielens = MovieLensDataset(
|
|
name="ml-1m", valid_ratio=0.2, test_ratio=0.1, verbose=True
|
|
)
|
|
g = movielens[0]
|
|
assert g.num_edges("user-movie") == g.num_edges("movie-user") == 1000209
|
|
|
|
movielens = MovieLensDataset(
|
|
name="ml-10m", valid_ratio=0.2, test_ratio=0.1, verbose=True
|
|
)
|
|
g = movielens[0]
|
|
assert g.num_edges("user-movie") == g.num_edges("movie-user") == 10000054
|