File size: 573 Bytes
a3682cf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 | import pickle
import time
from src.tgn.train import train_tgn
from src.tgn.evaluate import evaluate
def main():
start = time.time()
with open("data/graph/graph.pkl", "rb") as f:
graph_data = pickle.load(f)
model, memory, norm_stats = train_tgn(graph_data)
end = time.time()
print("Training complete")
print(f"Total runtime: {end - start:.2f} seconds")
roc, pr, probs, y_true = evaluate(model, memory, graph_data, norm_stats)
print(f"ROC-AUC: {roc:.4f}")
print(f"PR-AUC: {pr:.4f}")
if __name__ == "__main__":
main() |