search

ai4co / rl4co

A PyTorch library for all things Reinforcement Learning (RL) for Combinatorial Optimization (CO)
https://rl4.co
MIT License
455 stars 84 forks source link

[Feature enhancement] Edge handling in GCNEnc #60

Closed Junyoungpark closed 1 year ago

Junyoungpark commented 1 year ago

I found that the current implementation of GCNEnc cannot support the cases where we test the model that is not trained. This could be critical especially when we measure the generalization performance of the trained model.

An easy fix is, inside of the forward, when the given x.shape[1] != self.edge_index[0].max() - or something more smart -, we can reconstruct the edges from the inputs.

https://github.com/kaist-silab/rl4co/blob/daf558a8ae42a4c64ba8464135c3af85423ad469/rl4co/models/nn/graph/gcn.py#LL52C39-L52C49

cbhua commented 1 year ago

Good idea, added this feature in 1c22b6e64759ce0377baa28b8a471e2f714cd8b3.