def test_RGCN_apply_sparse_directed(): G, features = create_graph_features(is_directed=True) As = get_As(G) As = [A.tocoo() for A in As] A_indices = [ np.expand_dims(np.hstack((A.row[:, None], A.col[:, None])), 0) for A in As ] A_values = [np.expand_dims(A.data, 0) for A in As] generator = RelationalFullBatchNodeGenerator(G, sparse=True) rgcnModel = RGCN([2], generator, num_bases=10, activations=["relu"], dropout=0.5) x_in, x_out = rgcnModel.build() model = keras.Model(inputs=x_in, outputs=x_out) # Check fit method out_indices = np.array([[0, 1]], dtype="int32") preds_1 = model.predict([features[None, :, :], out_indices] + A_indices + A_values) assert preds_1.shape == (1, 2, 2) # Check fit_generator method preds_2 = model.predict_generator(generator.flow(["a", "b"])) assert preds_2.shape == (1, 2, 2) assert preds_1 == pytest.approx(preds_2)
def test_RGCN_apply_dense(): G, features = create_graph_features() As = get_As(G) As = [np.expand_dims(A.todense(), 0) for A in As] generator = RelationalFullBatchNodeGenerator(G, sparse=False) rgcnModel = RGCN([2], generator, num_bases=10, activations=["relu"], dropout=0.5) x_in, x_out = rgcnModel.build() model = keras.Model(inputs=x_in, outputs=x_out) # Check fit method out_indices = np.array([[0, 1]], dtype="int32") preds_1 = model.predict([features[None, :, :], out_indices] + As) assert preds_1.shape == (1, 2, 2) # Check fit_generator method preds_2 = model.predict_generator(generator.flow(["a", "b"])) assert preds_2.shape == (1, 2, 2) assert preds_1 == pytest.approx(preds_2)