Adding early stopping to examples and also fix staleness issue with DTDG model embeddings by shenyangHuang · Pull Request #397 · tgm-team/tgm

shenyangHuang · 2026-02-16T20:54:12Z

Summary / Description

Added early stopping to all linkproppred, nodeproppred and graphproppred examples

Related Issues: #394

Type of Change

Test Evidence

Requires running integration test to see if anything breaks

Questions / Discussion Points

For graphproppred, it is not clear what metric should use for early stopping @benjaminnNgo

…dtdg models

codecov · 2026-02-16T20:56:14Z

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

shenyangHuang · 2026-02-16T20:57:41Z


-test_results = eval(test_loader, test_labels, encoder, decoder, test_metrics)
-log_metrics_dict(test_results, epoch=args.epochs)
+    val_score = val_results['BinaryAUROC']


using Binary AUROC here as it is not clear which metric to use @benjaminnNgo

BinaryAUROC is good

Note: this fails on integration tests, KeyError for "BinaryAUROC": https://github.com/tgm-team/tgm/actions/runs/22492271186/job/65157420939

ntgbaoo

Fixing stale embeddings looks good to me.

For early stopping, I think we are missing "early stop" as the logic added in this PR just evaluates on the test set for the best epoch and does not really stop when over patience. Not sure if we want to support that.

For ICLR experiments, I have implemented the EarlyStopping class, and I remembered seeing @Jacob-Chmura implement something similar. Do we want to add this logic to the core? Or we can have utils.py within examples and have EarlyStopping class there.

ntgbaoo · 2026-02-17T03:48:16Z


-test_results = eval(test_loader, test_labels, encoder, decoder, test_metrics)
-log_metrics_dict(test_results, epoch=args.epochs)
+    val_score = val_results['BinaryAUROC']


BinaryAUROC is good

shenyangHuang · 2026-02-20T16:31:44Z

Fixing stale embeddings looks good to me.

For early stopping, I think we are missing "early stop" as the logic added in this PR just evaluates on the test set for the best epoch and does not really stop when over patience. Not sure if we want to support that.

For ICLR experiments, I have implemented the EarlyStopping class, and I remembered seeing @Jacob-Chmura implement something similar. Do we want to add this logic to the core? Or we can have utils.py within examples and have EarlyStopping class there.

I think it is a good idea to standardize it. I don't think early stopping should be in core as it doesn't relate to core TG workflow. we can discuss next week. For now, yeah this doesn't implement the full early stopping logic.

Jacob-Chmura

Running some more integration tests, but approving to unblock.

Jacob-Chmura · 2026-02-27T15:10:03Z

+    z = z.detach()
+    h_0 = h_0.detach()
+    c_0 = c_0.detach()


Not needed in eval?

maybe, but I think it is good practice to detach it to ensure no gradient will go into training the model from eval

Jacob-Chmura · 2026-02-27T15:10:38Z

                z = encoder(snapshot_batch, static_node_x)
            except StopIteration:
                pass
+    z = z.detach()


Similarly, not needed?

github-actions · 2026-02-27T15:19:29Z

Job: test_dygformer_nodeprop_pred_tgbn-trade

Metric	Value
discretize latency	0.004
materialize latency	0.002
execute_active_hooks latency	0.003
train latency	19.563
train peak_gpu_mb	524.286
train alloc_gpu_mb	519.450
eval latency	2.939
eval peak_gpu_mb	190.000
eval alloc_gpu_mb	167.660
Loss epoch 1	251.346
Validation ndcg epoch 1	0.278
Test ndcg epoch 1	0.228

Job: test_persistant_forecast_nodeprop_pred_tgbn-trade

Metric	Value
materialize latency	0.007
eval latency	0.244
Validation ndcg	0.860
Test ndcg	0.855

Job: test_gcn_nodeprop_pred_tgbn-trade

Metric	Value
materialize latency	0.004
train latency	0.309
train peak_gpu_mb	29.133
train alloc_gpu_mb	14.607
eval latency	0.068
eval peak_gpu_mb	29.108
eval alloc_gpu_mb	11.107
Loss epoch 1	112.194
Validation ndcg epoch 1	0.412
Test ndcg epoch 5	0.420
Loss epoch 2	92.109
Validation ndcg epoch 2	0.438
Loss epoch 3	89.943
Validation ndcg epoch 3	0.460
Loss epoch 4	88.273
Validation ndcg epoch 4	0.463
Loss epoch 5	87.408
Validation ndcg epoch 5	0.479

Job: test_tgcn_nodeprop_pred_tgbn-trade

Metric	Value
materialize latency	0.003
train latency	0.335
train peak_gpu_mb	40.831
train alloc_gpu_mb	24.871
eval latency	0.072
eval peak_gpu_mb	41.211
eval alloc_gpu_mb	21.574
Loss epoch 1	115.459
Validation ndcg epoch 1	0.363
Test ndcg epoch 5	0.353
Loss epoch 2	92.759
Validation ndcg epoch 2	0.369
Loss epoch 3	91.286
Validation ndcg epoch 3	0.361
Loss epoch 4	90.766
Validation ndcg epoch 4	0.394
Loss epoch 5	90.648
Validation ndcg epoch 5	0.374

Job: test_gclstm_nodeprop_pred_tgbn-trade

Metric	Value
materialize latency	0.003
train latency	0.351
train peak_gpu_mb	23.279
train alloc_gpu_mb	7.217
eval latency	0.072
eval peak_gpu_mb	22.988
eval alloc_gpu_mb	3.224
Loss epoch 1	111.865
Validation ndcg epoch 1	0.349
Test ndcg epoch 5	0.421
Loss epoch 2	89.703
Validation ndcg epoch 2	0.409
Loss epoch 3	88.076
Validation ndcg epoch 3	0.450
Loss epoch 4	85.389
Validation ndcg epoch 4	0.497
Loss epoch 5	81.936
Validation ndcg epoch 5	0.558

Job: test_tpnet_nodeprop_pred_tgbn-trade

Metric	Value
discretize latency	0.003
materialize latency	0.002
execute_active_hooks latency	0.003
train latency	17.095
train peak_gpu_mb	383.986
train alloc_gpu_mb	379.795
eval latency	2.648
eval peak_gpu_mb	364.905
eval alloc_gpu_mb	343.030
Loss epoch 1	252.953
Validation ndcg epoch 1	0.221
Test ndcg epoch 1	0.188

Job: test_tgat_nodeprop_pred_recency_sampler_tgbn-trade

Metric	Value
materialize latency	0.002
execute_active_hooks latency	0.000
train latency	5.064
train peak_gpu_mb	289.341
train alloc_gpu_mb	287.282
eval latency	0.582
eval peak_gpu_mb	155.127
eval alloc_gpu_mb	132.489
Loss epoch 1	242.634
Validation ndcg epoch 1	0.353
Test ndcg epoch 1	0.294

Download raw logs

Jacob-Chmura · 2026-02-27T21:31:26Z

tgm/examples/nodeproppred/tgn.py

Lines 215 to 230 in 72c8bf9

memory = TGNMemory(

full_data.num_nodes,

test_dg.edge_feats_dim,

args.memory_dim,

args.time_dim,

message_module=IdentityMessage(

test_dg.edge_feats_dim, args.memory_dim, args.time_dim

),

aggregator_module=LastAggregator(),

).to(args.device)

encoder = GraphAttentionEmbedding(

in_channels=args.memory_dim,

out_channels=args.embed_dim,

msg_dim=test_dg.edge_feats_dim,

time_enc=memory.time_enc,

).to(args.device)

unrelated to this PR, but I found that we forgot to update edge_feat_dim -> edge_x_dim. Let's fix this while we are updating this file, thanks

github-actions · 2026-03-06T19:49:03Z

Job: test_dygformer_nodeprop_pred_tgbn-trade

Metric	Value
discretize latency	0.004
materialize latency	0.002
execute_active_hooks latency	0.003
train latency	19.115
train peak_gpu_mb	524.286
train alloc_gpu_mb	519.450
eval latency	2.938
eval peak_gpu_mb	190.000
eval alloc_gpu_mb	167.660
Loss epoch 1	251.346
Validation ndcg epoch 1	0.278
Test ndcg epoch 1	0.228

Job: test_tgat_nodeprop_pred_recency_sampler_tgbn-trade

Metric	Value
materialize latency	0.002
execute_active_hooks latency	0.000
train latency	4.993
train peak_gpu_mb	289.341
train alloc_gpu_mb	287.282
eval latency	0.572
eval peak_gpu_mb	155.127
eval alloc_gpu_mb	132.489
Loss epoch 1	242.634
Validation ndcg epoch 1	0.353
Test ndcg epoch 1	0.294

Job: test_tgcn_nodeprop_pred_tgbn-trade

Metric	Value
materialize latency	0.003
train latency	0.251
train peak_gpu_mb	40.831
train alloc_gpu_mb	24.871
eval latency	0.072
eval peak_gpu_mb	41.211
eval alloc_gpu_mb	21.574
Loss epoch 1	115.459
Validation ndcg epoch 1	0.363
Test ndcg epoch 5	0.353
Loss epoch 2	92.759
Validation ndcg epoch 2	0.369
Loss epoch 3	91.286
Validation ndcg epoch 3	0.361
Loss epoch 4	90.766
Validation ndcg epoch 4	0.394
Loss epoch 5	90.648
Validation ndcg epoch 5	0.374

Job: test_gcn_nodeprop_pred_tgbn-trade

Metric	Value
materialize latency	0.003
train latency	0.273
train peak_gpu_mb	29.133
train alloc_gpu_mb	14.607
eval latency	0.064
eval peak_gpu_mb	29.108
eval alloc_gpu_mb	11.107
Loss epoch 1	112.194
Validation ndcg epoch 1	0.412
Test ndcg epoch 5	0.420
Loss epoch 2	92.109
Validation ndcg epoch 2	0.438
Loss epoch 3	89.943
Validation ndcg epoch 3	0.460
Loss epoch 4	88.273
Validation ndcg epoch 4	0.463
Loss epoch 5	87.408
Validation ndcg epoch 5	0.479

Job: test_tpnet_nodeprop_pred_tgbn-trade

Metric	Value
discretize latency	0.003
materialize latency	0.002
execute_active_hooks latency	0.003
train latency	17.657
train peak_gpu_mb	383.986
train alloc_gpu_mb	379.795
eval latency	2.687
eval peak_gpu_mb	364.905
eval alloc_gpu_mb	343.030
Loss epoch 1	252.953
Validation ndcg epoch 1	0.221
Test ndcg epoch 1	0.188

Job: test_gclstm_nodeprop_pred_tgbn-trade

Metric	Value
materialize latency	0.003
train latency	0.253
train peak_gpu_mb	23.279
train alloc_gpu_mb	7.217
eval latency	0.070
eval peak_gpu_mb	22.988
eval alloc_gpu_mb	3.224
Loss epoch 1	111.865
Validation ndcg epoch 1	0.349
Test ndcg epoch 5	0.421
Loss epoch 2	89.703
Validation ndcg epoch 2	0.409
Loss epoch 3	88.076
Validation ndcg epoch 3	0.450
Loss epoch 4	85.389
Validation ndcg epoch 4	0.497
Loss epoch 5	81.936
Validation ndcg epoch 5	0.558

Job: test_persistant_forecast_nodeprop_pred_tgbn-trade

Metric	Value
materialize latency	0.005
eval latency	0.254
Validation ndcg	0.860
Test ndcg	0.855

Download raw logs

Unblock

adding early stopping and also fixes last embedding stale issue with …

4edef00

…dtdg models

shenyangHuang added this to the Beta Release 3 milestone Feb 16, 2026

shenyangHuang requested review from Jacob-Chmura and ntgbaoo February 16, 2026 20:54

shenyangHuang self-assigned this Feb 16, 2026

shenyangHuang linked an issue Feb 16, 2026 that may be closed by this pull request

Adding Early Stopping to all examples #394

Closed

adding gcn example

d2c06d6

shenyangHuang commented Feb 16, 2026

View reviewed changes

ntgbaoo previously requested changes Feb 17, 2026

View reviewed changes

Jacob-Chmura approved these changes Feb 27, 2026

View reviewed changes

Jacob-Chmura reviewed Feb 27, 2026

View reviewed changes

changed edge_feat_dim to edge_x_dim

e9e56fb

shenyangHuang merged commit 0600303 into main Mar 20, 2026
7 of 8 checks passed

	memory = TGNMemory(
	full_data.num_nodes,
	test_dg.edge_feats_dim,
	args.memory_dim,
	args.time_dim,
	message_module=IdentityMessage(
	test_dg.edge_feats_dim, args.memory_dim, args.time_dim
	),
	aggregator_module=LastAggregator(),
	).to(args.device)
	encoder = GraphAttentionEmbedding(
	in_channels=args.memory_dim,
	out_channels=args.embed_dim,
	msg_dim=test_dg.edge_feats_dim,
	time_enc=memory.time_enc,
	).to(args.device)

Uh oh!

Conversation

shenyangHuang commented Feb 16, 2026

Summary / Description

Type of Change

Test Evidence

Questions / Discussion Points

Uh oh!

codecov Bot commented Feb 16, 2026

Codecov Report

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ntgbaoo left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shenyangHuang commented Feb 20, 2026

Uh oh!

Jacob-Chmura left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Feb 27, 2026

Job: test_dygformer_nodeprop_pred_tgbn-trade

Job: test_persistant_forecast_nodeprop_pred_tgbn-trade

Job: test_gcn_nodeprop_pred_tgbn-trade

Job: test_tgcn_nodeprop_pred_tgbn-trade

Job: test_gclstm_nodeprop_pred_tgbn-trade

Job: test_tpnet_nodeprop_pred_tgbn-trade

Job: test_tgat_nodeprop_pred_recency_sampler_tgbn-trade

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Mar 6, 2026

Job: test_dygformer_nodeprop_pred_tgbn-trade

Job: test_tgat_nodeprop_pred_recency_sampler_tgbn-trade

Job: test_tgcn_nodeprop_pred_tgbn-trade

Job: test_gcn_nodeprop_pred_tgbn-trade

Job: test_tpnet_nodeprop_pred_tgbn-trade

Job: test_gclstm_nodeprop_pred_tgbn-trade

Job: test_persistant_forecast_nodeprop_pred_tgbn-trade

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants