import os
from tqdm import tqdm
import typing
from datetime import datetime
from time import time
import mgp
import torch
from torch_geometric.nn import to_hetero
from torch_geometric.data import HeteroData
from mage.node_classification.models.gat import GAT
from mage.node_classification.models.gatv2 import GATv2
from mage.node_classification.models.sage import SAGE
from mage.node_classification.models.gatjk import GATJK
from mage.node_classification.utils.metrics import metrics
from mage.node_classification.utils.extract_from_database import extract_from_database
from mage.node_classification.models.train_model import train_epoch
##############################
# constants
##############################
# parameters for the model
class ModelParams:
IN_CHANNELS = "in_channels"
OUT_CHANNELS = "out_channels"
HIDDEN_FEATURES_SIZE = "hidden_features_size"
LAYER_TYPE = "layer_type"
AGGREGATOR = "aggregator"
# parameters for optimizer
class OptimizerParams:
LEARNING_RATE = "learning_rate"
WEIGHT_DECAY = "weight_decay"
# parameters for data
class DataParams:
SPLIT_RATIO = "split_ratio"
METRICS = "metrics"
# parameters relevant to memgraph database
class MemgraphParams:
NODE_ID_PROPERTY = "node_id_property"
# parameters for training
class TrainParams:
NUM_EPOCHS = "num_epochs"
CONSOLE_LOG_FREQ = "console_log_freq"
CHECKPOINT_FREQ = "checkpoint_freq"
BATCH_SIZE = "batch_size"
MAX_MODELS_TO_KEEP = "max_models_to_keep"
TIME_BETWEEN_CHECKPOINTS = "time_between_checkpoints"
# parameters relevant for heterogeneous structure
class HeteroParams:
FEATURES_NAME = "features_name"
OBSERVED_ATTRIBUTE = "observed_attribute"
CLASS_NAME = "class_name"
REINDEXING = "reindexing"
INV_REINDEXING = "inv_reindexing"
NUM_NODES_SAMPLE = "num_nodes_sample"
NUM_ITERATIONS_SAMPLE = "num_iterations_sample"
LABEL_REINDEXING = "label_reindexing"
INV_LABEL_REINDEXING = "inv_label_reindexing"
# other necessary parameters
class OtherParams:
DEVICE_TYPE = "device_type"
PATH_TO_MODEL = "path_to_model"
PATIENCE = "patience"
MODEL_SAVING_FOLDER = "model_saving_folder"
GAT_MODEL = "GAT"
GATV2_MODEL = "GATv2"
SAGE_MODEL = "SAGE"
GAT_WITH_JK = "GATJK"
# dictionary of models
MODELS = {GAT_MODEL: GAT, GATV2_MODEL: GATv2, SAGE_MODEL: SAGE, GAT_WITH_JK: GATJK}
global model, current_values
model: mgp.Any = None
current_values: typing.Dict = {}
# list for saving logged data
logged_data: mgp.List = []
# dictionary of defined input types
DEFINED_INPUT_TYPES = {
ModelParams.HIDDEN_FEATURES_SIZE: list,
ModelParams.LAYER_TYPE: str,
TrainParams.NUM_EPOCHS: int,
OptimizerParams.LEARNING_RATE: float,
OptimizerParams.WEIGHT_DECAY: float,
DataParams.SPLIT_RATIO: float,
MemgraphParams.NODE_ID_PROPERTY: str,
OtherParams.DEVICE_TYPE: str,
TrainParams.CONSOLE_LOG_FREQ: int,
TrainParams.CHECKPOINT_FREQ: int,
TrainParams.BATCH_SIZE: int,
TrainParams.MAX_MODELS_TO_KEEP: int,
TrainParams.TIME_BETWEEN_CHECKPOINTS: float,
ModelParams.AGGREGATOR: str,
DataParams.METRICS: list,
HeteroParams.OBSERVED_ATTRIBUTE: str,
HeteroParams.FEATURES_NAME: str,
HeteroParams.CLASS_NAME: str,
HeteroParams.REINDEXING: dict,
HeteroParams.INV_REINDEXING: dict,
HeteroParams.NUM_NODES_SAMPLE: int,
HeteroParams.NUM_ITERATIONS_SAMPLE: int,
OtherParams.PATH_TO_MODEL: str,
OtherParams.PATIENCE: int,
OtherParams.MODEL_SAVING_FOLDER: str,
}
# dictionary of default values for input types
DEFAULT_VALUES = {
ModelParams.HIDDEN_FEATURES_SIZE: [16, 16],
ModelParams.LAYER_TYPE: "GATJK",
TrainParams.NUM_EPOCHS: 100,
OptimizerParams.LEARNING_RATE: 0.1,
OptimizerParams.WEIGHT_DECAY: 5e-4,
DataParams.SPLIT_RATIO: 0.8,
MemgraphParams.NODE_ID_PROPERTY: "id",
OtherParams.DEVICE_TYPE: "cpu",
TrainParams.CONSOLE_LOG_FREQ: 5,
TrainParams.CHECKPOINT_FREQ: 5,
TrainParams.BATCH_SIZE: 64,
TrainParams.MAX_MODELS_TO_KEEP: 5,
TrainParams.TIME_BETWEEN_CHECKPOINTS: 2.0,
ModelParams.AGGREGATOR: "mean",
DataParams.METRICS: [
"loss",
"accuracy",
"f1_score",
"precision",
"recall",
"num_wrong_examples",
],
HeteroParams.OBSERVED_ATTRIBUTE: "",
HeteroParams.FEATURES_NAME: "features",
HeteroParams.CLASS_NAME: "class",
HeteroParams.REINDEXING: {},
HeteroParams.INV_REINDEXING: {},
HeteroParams.NUM_NODES_SAMPLE: 512,
HeteroParams.NUM_ITERATIONS_SAMPLE: 4,
OtherParams.PATH_TO_MODEL: "",
OtherParams.PATIENCE: 10,
OtherParams.MODEL_SAVING_FOLDER: "/tmp/torch_models",
}
##############################
# set model parameters
##############################
def declare_data(ctx: mgp.ProcCtx) -> HeteroData:
"""This function initializes global variable data.
Args:
ctx (mgp.ProcCtx): current context
"""
global current_values
# change device type to cuda if possible
current_values[OtherParams.DEVICE_TYPE] = torch.device(
"cuda:0" if torch.cuda.is_available() else "cpu"
)
nodes = list(iter(ctx.graph.vertices)) # obtain nodes from context
if not nodes:
raise Exception("Graph is empty.")
# extraction of data from database to torch.Tensors
(
data,
current_values[HeteroParams.OBSERVED_ATTRIBUTE],
current_values[HeteroParams.REINDEXING],
current_values[HeteroParams.INV_REINDEXING],
current_values[HeteroParams.LABEL_REINDEXING],
current_values[HeteroParams.INV_LABEL_REINDEXING],
) = extract_from_database(
nodes,
current_values[DataParams.SPLIT_RATIO],
current_values[HeteroParams.FEATURES_NAME],
current_values[HeteroParams.CLASS_NAME],
current_values[OtherParams.DEVICE_TYPE],
)
observed_attribute_data = data[current_values[HeteroParams.OBSERVED_ATTRIBUTE]]
# second parameter of shape of feature matrix is number of input channels
current_values[ModelParams.IN_CHANNELS] = observed_attribute_data.x.size(dim=1)
# number of output channels is number of classes in the dataset
current_values[ModelParams.OUT_CHANNELS] = len(
set(observed_attribute_data.y.detach().cpu().numpy())
)
return data
def declare_model(data: mgp.Any):
"""This function initializes global variables model, opt and criterion.
Args:
ctx (mgp.ProcCtx): current context
"""
# choose one of the available layer types
global model, current_values
args_gatjk = [
current_values[ModelParams.IN_CHANNELS],
current_values[ModelParams.HIDDEN_FEATURES_SIZE],
current_values[ModelParams.OUT_CHANNELS],
]
args_inductive = [
current_values[ModelParams.IN_CHANNELS],
current_values[ModelParams.HIDDEN_FEATURES_SIZE],
current_values[ModelParams.OUT_CHANNELS],
current_values[ModelParams.AGGREGATOR],
]
# choose model architecture according to layer type
layer_type = current_values[ModelParams.LAYER_TYPE]
if layer_type not in MODELS.keys():
raise Exception(
"You didn't choose one of currently available models (GAT, GATv2, GATJK and SAGE). Please choose one of them."
)
args = args_gatjk if layer_type == GAT_WITH_JK else args_inductive
model = MODELS[layer_type](*args)
# convert model to hetero structure
# (if graph is homogeneous, we also do this conversion since all calculations are same)
metadata = (data.node_types, data.edge_types)
model = to_hetero(model, metadata)
# move model to device
model.to(current_values[OtherParams.DEVICE_TYPE])
# set default optimizer
opt = torch.optim.Adam(
model.parameters(),
lr=current_values[OptimizerParams.LEARNING_RATE],
weight_decay=current_values[OptimizerParams.WEIGHT_DECAY],
)
# set default criterion
criterion = torch.nn.CrossEntropyLoss()
return opt, criterion
def declare_saving_paths():
"""This function initializes global variables paths."""
global current_values
# either make new folder for saving models, or use existing one with exactly this name
try:
path = os.path.join(
os.getcwd(), current_values[OtherParams.MODEL_SAVING_FOLDER]
)
os.makedirs(path)
print(f"New folder for saving models was created on destination {path}.")
except FileExistsError:
print(f"Folder for saving models already exists on destination {path}.")
current_values[OtherParams.PATH_TO_MODEL] = os.path.join(
os.getcwd(),
current_values[OtherParams.MODEL_SAVING_FOLDER],
"model_" + current_values[ModelParams.LAYER_TYPE] + "_",
)
@mgp.read_proc
def set_model_parameters(
params: mgp.Any = {},
) -> mgp.Record(
hidden_features_size=list,
layer_type=str,
aggregator=str,
learning_rate=float,
weight_decay=float,
split_ratio=float,
metrics=mgp.Any,
node_id_property=str,
num_epochs=int,
console_log_freq=int,
checkpoint_freq=int,
device_type=str,
path_to_model=str,
):
"""The purpose of this function is to initialize all global variables.
_You_ can change those via **params** dictionary.
It checks if variables in **params** are defined appropriately. If so,
map of default global parameters is overridden with user defined dictionary params.
After that it executes previously defined functions declare_globals and
declare_model_and_data and sets each global variable to some value.
Args:
ctx: (mgp.ProcCtx): current context,
params: (mgp.Map, optional): user defined parameters from query module. Defaults to {}
Raises:
Exception: exception is raised if some variable in dictionary params is not
defined as it should be
Returns:
mgp.Record(
hidden_features_size (list): list of hidden features
layer_type (str): type of layer
aggregator (str): type of aggregator
learning_rate (float): learning rate
weight_decay (float): weight decay
split_ratio (float): ratio between training and validation data
metrics (list): list of metrics to be calculated
node_id_property (str): name of nodes id property
num_epochs (int): number of epochs
console_log_freq (int): frequency of logging metrics
checkpoint_freq (int): frequency of saving models
device_type (str): cpu or cuda
path_to_model (str): path where model is load and saved
)
"""
global DEFINED_INPUT_TYPES, DEFAULT_VALUES, current_values
# function checks if input values in dictionary are correctly typed
def is_correctly_typed(defined_types, input_values):
if isinstance(defined_types, dict) and isinstance(input_values, dict):
# defined_types is a dict of types
return all(
k in input_values # check if exists
and is_correctly_typed(
defined_types[k], input_values[k]
) # check for correct type
for k in defined_types
)
elif isinstance(defined_types, type):
return isinstance(input_values, defined_types)
else:
return False
# hidden_features_size and metrics are sometimes translated as tuples,
# which are not hashable, but conversion to lists makes them hashable
if (
ModelParams.HIDDEN_FEATURES_SIZE in params.keys()
and isinstance(params[ModelParams.HIDDEN_FEATURES_SIZE], tuple)
):
params[ModelParams.HIDDEN_FEATURES_SIZE] = list(
params[ModelParams.HIDDEN_FEATURES_SIZE]
)
if (
DataParams.METRICS in params.keys()
and isinstance(params[DataParams.METRICS], tuple)
):
params[DataParams.METRICS] = list(params[DataParams.METRICS])
# override any default parameters
current_values = {**DEFAULT_VALUES, **params}
# raise exception if some variable in dictionary params is not defined as it should be
if not is_correctly_typed(DEFINED_INPUT_TYPES, current_values):
raise Exception("Input dictionary is not correctly typed.")
# define paths
declare_saving_paths()
return mgp.Record(
hidden_features_size=current_values[ModelParams.HIDDEN_FEATURES_SIZE],
layer_type=current_values[ModelParams.LAYER_TYPE],
aggregator=current_values[ModelParams.AGGREGATOR],
learning_rate=current_values[OptimizerParams.LEARNING_RATE],
weight_decay=current_values[OptimizerParams.WEIGHT_DECAY],
split_ratio=current_values[DataParams.SPLIT_RATIO],
metrics=current_values[DataParams.METRICS],
node_id_property=current_values[MemgraphParams.NODE_ID_PROPERTY],
num_epochs=current_values[TrainParams.NUM_EPOCHS],
console_log_freq=current_values[TrainParams.CONSOLE_LOG_FREQ],
checkpoint_freq=current_values[TrainParams.CHECKPOINT_FREQ],
device_type=current_values[OtherParams.DEVICE_TYPE],
path_to_model=current_values[OtherParams.PATH_TO_MODEL],
)
##############################
# train
##############################
def fetch_saved_models():
"""The purpose of this function is to fetch all saved models.
Returns:
model_saving_folder (str): path to folder with saved models
models (list): list of paths of saved models
"""
global model
model_saving_folder = os.path.join(current_values[OtherParams.MODEL_SAVING_FOLDER])
models = [
f
for f in os.listdir(model_saving_folder)
if os.path.isfile(os.path.join(model_saving_folder, f))
and f.endswith(".pt")
and f.startswith("model")
]
models.sort(reverse=True)
return model_saving_folder, models
def save_model_to_folder() -> str:
"""The purpose of this function is to save model to folder.
Returns:
path_to_saved_model (str): path to saved model
"""
model_saving_folder, models = fetch_saved_models()
# delete oldest models if there are more than max models to keep
for i in range(current_values[TrainParams.MAX_MODELS_TO_KEEP] - 1, len(models)):
os.remove(os.path.join(model_saving_folder, models[i]))
path_to_saved_model = (
current_values[OtherParams.PATH_TO_MODEL]
+ datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
+ ".pt"
)
torch.save(
model.state_dict(),
path_to_saved_model,
)
return path_to_saved_model
@mgp.read_proc
def train(
ctx: mgp.ProcCtx, num_epochs: int = 100
) -> mgp.Record(
epoch=int, loss=float, val_loss=float, train_log=mgp.Any, val_log=mgp.Any
):
"""This function performs training of model. It first declares data, model,
optimizer and criterion. Then it performs training.
Args:
ctx (mgp.ProcCtx): context of process
num_epochs (int, optional): number of epochs. Defaults to 100.
Raises:
Exception: raised if graph is empty
Returns:
list of mgp.Record of
epoch (int): epoch number
loss (float): loss of model on training data
val_loss (float): loss of model on validation data
train_log (list): list of metrics on training data
val_log (list): list of metrics on validation data
"""
global model, current_values, logged_data
# define fresh data
data = declare_data(ctx)
# define model
opt, criterion = declare_model(data)
current_values[TrainParams.NUM_EPOCHS] = num_epochs
num_nodes_sample = current_values[HeteroParams.NUM_NODES_SAMPLE]
num_iterations_sample = current_values[HeteroParams.NUM_ITERATIONS_SAMPLE]
# variables for early stopping
last_loss = float("inf")
trigger_times = 0
last_time = time()
# training
for epoch in tqdm(range(1, num_epochs + 1)):
# one epoch of training, both training and validation loss are returned
loss, val_loss = train_epoch(
model,
opt,
data,
criterion,
current_values[TrainParams.BATCH_SIZE],
current_values[HeteroParams.OBSERVED_ATTRIBUTE],
{
key: [num_nodes_sample] * num_iterations_sample
for key in data.node_types
},
)
# early stopping
if val_loss > last_loss:
trigger_times += 1
drop_epochs = (
str(trigger_times)
+ " "
+ ("consecutive epochs" if trigger_times > 1 else "consecutive epoch")
)
times_until_stopping = current_values[OtherParams.PATIENCE] - trigger_times
stop_after = (
str(times_until_stopping)
+ " "
+ ("more drops" if times_until_stopping > 1 else "more drop")
)
print(f"Loss has dropped for {drop_epochs}. Stopping after {stop_after}.")
if trigger_times >= current_values[OtherParams.PATIENCE]:
print("Early stopping!")
break
else:
trigger_times = 0
last_loss = val_loss
# log data every console_log_freq epochs
if epoch % current_values[TrainParams.CONSOLE_LOG_FREQ] == 0:
model.eval()
out = model(data.x_dict, data.edge_index_dict)
dict_train = metrics(
data[current_values[HeteroParams.OBSERVED_ATTRIBUTE]].train_mask,
out,
data,
current_values[DataParams.METRICS],
current_values[HeteroParams.OBSERVED_ATTRIBUTE],
current_values[OtherParams.DEVICE_TYPE],
)
dict_val = metrics(
data[current_values[HeteroParams.OBSERVED_ATTRIBUTE]].val_mask,
out,
data,
current_values[DataParams.METRICS],
current_values[HeteroParams.OBSERVED_ATTRIBUTE],
current_values[OtherParams.DEVICE_TYPE],
)
logged_data.append(
{
"epoch": epoch,
"loss": loss,
"val_loss": val_loss,
"train": dict_train,
"val": dict_val,
}
)
print(
f"Epoch: {epoch:03d}, Loss: {loss:.4f}, Val Loss: {val_loss:.4f},"
+ f'Accuracy: {logged_data[-1]["train"]["accuracy"]:.4f}, Accuracy: {logged_data[-1]["val"]["accuracy"]:.4f}'
)
# save model every checkpoint_freq epochs
if epoch % current_values[TrainParams.CHECKPOINT_FREQ] == 0:
if (
time() - last_time
> current_values[TrainParams.TIME_BETWEEN_CHECKPOINTS]
):
save_model_to_folder()
last_time = time()
return [
mgp.Record(
epoch=data["epoch"],
loss=data["loss"],
val_loss=data["val_loss"],
train_log=data["train"],
val_log=data["val"],
)
for data in logged_data
]
##############################
# get training data
##############################
@mgp.read_proc
def get_training_data() -> (
mgp.Record(
epoch=int, loss=float, val_loss=float, train_log=mgp.Any, val_log=mgp.Any
)
):
"""This function is used so user can see what is logged data from training.
Returns:
mgp.Record(
epoch (int): epoch number of record of logged data row
loss (float): loss in logged data row
val_loss (float): validation loss in logged data row
train_log (mgp.Any): training parameters of record of logged data row
val_log (mgp.Any): validation parameters of record of logged data row
): record to return
"""
return [
mgp.Record(
epoch=data["epoch"],
loss=data["loss"],
val_loss=data["val_loss"],
train_log=data["train"],
val_log=data["val"],
)
for data in logged_data
]
##############################
# model loading and saving, predict
##############################
@mgp.read_proc
def save_model() -> mgp.Record(path=str, status=str):
"""This function saves model to model saving folder. If there are already total
of max_models_to_keep models in model saving folder, oldest model is deleted.
Exception: raised if model is not initialized or defined
Returns:
mgp.Record(
path (str): path to saved model
status (str): status of saving model
): return record
"""
if model is None:
raise Exception(
"There are no initialized or loaded models. First load or initialize a model to be able save it."
)
path_to_saved_model = save_model_to_folder()
return mgp.Record(
path=path_to_saved_model, status="Model has been successfully saved."
)
@mgp.read_proc
def load_model(ctx: mgp.ProcCtx, num: int = 0) -> mgp.Record(path=str, status=str):
"""This function loads model from defined folder for saved models.
Args:
num (int, optional): ordinary number of model to load from default map. Defaults to 0 (newest model).
Returns:
mgp.Record(path (str): path to loaded model): return record
"""
global model
data = declare_data(ctx)
declare_model(data)
model_saving_folder, models = fetch_saved_models()
if len(models) == 0:
raise Exception("There are no saved models.")
if len(models) < (len(models) + num) % len(models) + 1:
raise Exception(
f"Model with number {num} does not exist. There are {len(models)} models saved."
)
path_to_load_model = os.path.join(model_saving_folder, models[num])
model.load_state_dict(torch.load(path_to_load_model))
return mgp.Record(
path=path_to_load_model, status="Model has been successfully loaded."
)
@mgp.read_proc
def predict(
ctx: mgp.ProcCtx, vertex: mgp.Vertex
) -> mgp.Record(predicted_class=int, status=str):
"""This function predicts metrics on one node. It is suggested that user previously
loads unseen test data to predict on it.
Example of usage:
MATCH (n {id: 1}) CALL node_classification.predict(n) YIELD * RETURN predicted_class;
# note: if node with property id = 1 doesn't exist, query module won't be called
Args:
ctx (mgp.ProcCtx): proc context
vertex (mgp.Vertex): node to predict on
Returns:
mgp.Record(
predicted_class (int): predicted class
status (str): status of prediction
): record to return
"""
global current_values
# define fresh data
data = declare_data(ctx)
if model is None:
raise Exception("Load a model before predicting.")
model.eval()
out = model(data.x_dict, data.edge_index_dict)
pred = out[current_values[HeteroParams.OBSERVED_ATTRIBUTE]].argmax(dim=1)
inv_reindexing = HeteroParams.INV_REINDEXING
observed_attribute = current_values[HeteroParams.OBSERVED_ATTRIBUTE]
position = current_values[inv_reindexing][observed_attribute][vertex.id]
predicted_class = int(pred.detach().cpu().numpy()[position])
return mgp.Record(
predicted_class=current_values[HeteroParams.INV_LABEL_REINDEXING][
predicted_class
],
status="Prediction complete.",
)
@mgp.read_proc
def reset() -> mgp.Record(status=str):
"""This function resets all variables to default values.
Returns:
mgp.Record(status (str): status of reset): record to return
"""
# set model and logged_data to None
global model, current_values, logged_data
model = None
logged_data = []
# reinitialize current_values
current_values = DEFAULT_VALUES
return mgp.Record(status="Global parameters and logged data have been reset")