Model objects from the model catalog

Models of the GDS Model Catalog are represented as Model objects in the Python client, similar to how there are graph objects. Model objects are typically constructed from training a pipeline or a GraphSAGE model, in which case a reference to the trained model in the form of a Model object is returned.

Once created, the Model objects can be passed as arguments to methods in the Python client, such as the model catalog operations. Additionally, the Model objects have convenience methods allowing for inspection of the models represented without explicitly involving the model catalog.

In the examples below we assume that we have an instantiated GraphDataScience object called gds. Read more about this in Getting started.

Constructing a model object

The primary way to construct a model object is through training a model. There are two types of models: pipeline models and GraphSAGE models. In order to train a pipeline model, a pipeline must first be created and configured Read more about how to operate pipelines in Machine learning pipelines, including examples of using pipeline models. In this section, we will exemplify creating and using a GraphSAGE model object.

First, we introduce a small road-network graph:

gds.run_cypher(
  """
  CREATE
    (a:City {name: "New York City", settled: 1624}),
    (b:City {name: "Philadelphia", settled: 1682}),
    (c:City:Capital {name: "Washington D.C.", settled: 1790}),
    (d:City {name: "Baltimore", settled: 1729}),
    (e:City {name: "Atlantic City", settled: 1854}),
    (f:City {name: "Boston", settled: 1822}),

    (a)-[:ROAD {cost: 50}]->(b),
    (a)-[:ROAD {cost: 50}]->(c),
    (a)-[:ROAD {cost: 100}]->(d),
    (b)-[:ROAD {cost: 40}]->(d),
    (c)-[:ROAD {cost: 40}]->(d),
    (c)-[:ROAD {cost: 80}]->(e),
    (d)-[:ROAD {cost: 30}]->(e),
    (d)-[:ROAD {cost: 80}]->(f),
    (e)-[:ROAD {cost: 40}]->(f);
  """
)
G, project_result = gds.graph.project.native(
    "road_graph",
    "City",
    "ROAD",
    node_properties=["settled"],
    relationship_properties=["cost"]
)

assert G.relationship_count() == 9

Now we can use the graph G to train a GraphSage model.

model, train_result = gds.graph_sage.train(G, model_name="city-representation", feature_properties=["settled"], random_seed=42, concurrency=4)

assert train_result.model_info["metrics"]["ranEpochs"] == 1

where model is the model object, and train_result is a GraphSageTrainResult object containing metadata from the underlying procedure call.

Similarly, we can also get model objects from training machine learning pipelines.

To inspect a model that is already present in the model catalog, call the get method with its name, which returns a Models object:

model = gds.graph_sage.get("city-representation")

assert model.name() == "city-representation"

Inspecting a model object

All model objects expose convenience methods to inspect the represented model (name, exists, details) and to operate on it in the model catalog (drop, and, on GDS Enterprise Edition, store, load, publish, delete). See the Model API reference for the full list of methods and their signatures.

The details() method returns a ModelDetails object.

For example, to get the train configuration of our model object model created above, we would do the following:

train_config = model.details().train_config

assert train_config["concurrency"] == 4

Using a model object

The primary way to use model objects is for prediction. How to do so for GraphSAGE is described below, and on the Machine learning pipelines page for pipelines.

Additionally, model objects can be used as input to GDS Model Catalog operations. For instance, supposing we have our model object model created above, we could:

# Store the model on disk (GDS Enterprise Edition)
_ = model.store()

model.drop()

# Load the model again for further use
model.load()

GraphSAGE

As exemplified above in Constructing a model object, training a GraphSAGE model with the Python client is analogous to its Cypher counterpart.

Once trained, in addition to the methods above, the GraphSAGE model object will have the following methods.

predict_stream — predict embeddings for nodes of the input graph and stream the results.
predict_mutate — predict embeddings for nodes of the input graph and mutate the graph with the results.
predict_write — predict embeddings for nodes of the input graph and write the results back to the database.
predict_estimate — estimate the memory needed to predict embeddings for the input graph.

See the GraphSageModel API reference for the full method signatures.

The training metrics are available via model.details().model_info["metrics"].

So given the GraphSAGE model model we trained above, we could do the following:

# Make sure our training actually converged
metrics = model.details().model_info["metrics"]
assert metrics["didConverge"], "did not converge"

# Predict on `G` and write embedding node properties back to the database
predict_result = model.predict_write(G, write_property="embedding")
assert predict_result.node_properties_written == G.node_count()