mirror of
https://github.com/NicolasBohn/NexQuant.git
synced 2026-07-27 23:47:46 +00:00
Align factor coder into new framework (#47)
* use CoSTEER as component name * rename factorimplementation to avoid confusion * rename modelimplementation * align benchmark and evolving evaluators * add scenario to evaluator init function * rename all factorimplementationknowledge in CoSTEER * remove all scenario related information in component * remove useless code --------- Co-authored-by: xuyang1 <xuyang1@microsoft.com>
This commit is contained in:
Xu Yang
@@ -0,0 +1,72 @@
|
||||
# TODO: inherent from the benchmark base class
|
||||
import torch
|
||||
|
||||
from rdagent.components.coder.model_coder.model import ModelImplementation
|
||||
|
||||
|
||||
def get_data_conf(init_val):
|
||||
# TODO: design this step in the workflow
|
||||
in_dim = 1000
|
||||
in_channels = 128
|
||||
exec_config = {"model_eval_param_init": init_val}
|
||||
node_feature = torch.randn(in_dim, in_channels)
|
||||
edge_index = torch.randint(0, in_dim, (2, 2000))
|
||||
return (node_feature, edge_index), exec_config
|
||||
|
||||
|
||||
class ModelImpValEval:
|
||||
"""
|
||||
Evaluate the similarity of the model structure by changing the input and observe the output.
|
||||
|
||||
Assumption:
|
||||
- If the model structure is similar, the output will change in similar way when we change the input.
|
||||
|
||||
Challenge:
|
||||
- The key difference between it and implementing factors is that we have parameters in the layers (Factor operators often have no parameters or are given parameters).
|
||||
- we try to initialize the model param in similar value. So only the model structure is different.
|
||||
|
||||
Comparing the correlation of following sequences
|
||||
- modelA[init1](input1).hidden_out1, modelA[init1](input2).hidden_out1, ...
|
||||
- modelB[init1](input1).hidden_out1, modelB[init1](input2).hidden_out1, ...
|
||||
|
||||
For each hidden output, we can calculate a correlation. The average correlation will be the metrics.
|
||||
"""
|
||||
|
||||
def evaluate(self, gt: ModelImplementation, gen: ModelImplementation):
|
||||
round_n = 10
|
||||
|
||||
eval_pairs: list[tuple] = []
|
||||
|
||||
# run different input value
|
||||
for _ in range(round_n):
|
||||
# run different model initial parameters.
|
||||
for init_val in [-0.2, -0.1, 0.1, 0.2]:
|
||||
data, exec_config = get_data_conf(init_val)
|
||||
gt_res = gt.execute(data=data, config=exec_config)
|
||||
res = gen.execute(data=data, config=exec_config)
|
||||
eval_pairs.append((res, gt_res))
|
||||
|
||||
# flat and concat the output
|
||||
res_batch, gt_res_batch = [], []
|
||||
for res, gt_res in eval_pairs:
|
||||
res_batch.append(res.reshape(-1))
|
||||
gt_res_batch.append(gt_res.reshape(-1))
|
||||
res_batch = torch.stack(res_batch)
|
||||
gt_res_batch = torch.stack(gt_res_batch)
|
||||
|
||||
res_batch = res_batch.detach().numpy()
|
||||
gt_res_batch = gt_res_batch.detach().numpy()
|
||||
|
||||
# pearson correlation of each hidden output
|
||||
def norm(x):
|
||||
return (x - x.mean(axis=0)) / x.std(axis=0)
|
||||
|
||||
dim_corr = (norm(res_batch) * norm(gt_res_batch)).mean(axis=0) # the correlation of each hidden output
|
||||
|
||||
# aggregate all the correlation
|
||||
avr_corr = dim_corr.mean()
|
||||
# FIXME:
|
||||
# It is too high(e.g. 0.944) .
|
||||
# Check if it is not a good evaluation!!
|
||||
# Maybe all the same initial params will results in extreamly high correlation without regard to the model structure.
|
||||
return avr_corr
|
||||
@@ -0,0 +1,134 @@
|
||||
import math
|
||||
from typing import Any, Callable, Dict, Optional, Union
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch.nn import Parameter
|
||||
from torch_geometric.nn.conv import GCNConv, MessagePassing
|
||||
from torch_geometric.nn.inits import zeros
|
||||
from torch_geometric.nn.resolver import activation_resolver
|
||||
from torch_geometric.typing import Adj
|
||||
|
||||
|
||||
class AntiSymmetricConv(torch.nn.Module):
|
||||
r"""The anti-symmetric graph convolutional operator from the
|
||||
`"Anti-Symmetric DGN: a stable architecture for Deep Graph Networks"
|
||||
<https://openreview.net/forum?id=J3Y7cgZOOS>`_ paper.
|
||||
|
||||
.. math::
|
||||
\mathbf{x}^{\prime}_i = \mathbf{x}_i + \epsilon \cdot \sigma \left(
|
||||
(\mathbf{W}-\mathbf{W}^T-\gamma \mathbf{I}) \mathbf{x}_i +
|
||||
\Phi(\mathbf{X}, \mathcal{N}_i) + \mathbf{b}\right),
|
||||
|
||||
where :math:`\Phi(\mathbf{X}, \mathcal{N}_i)` denotes a
|
||||
:class:`~torch.nn.conv.MessagePassing` layer.
|
||||
|
||||
Args:
|
||||
in_channels (int): Size of each input sample.
|
||||
phi (MessagePassing, optional): The message passing module
|
||||
:math:`\Phi`. If set to :obj:`None`, will use a
|
||||
:class:`~torch_geometric.nn.conv.GCNConv` layer as default.
|
||||
(default: :obj:`None`)
|
||||
num_iters (int, optional): The number of times the anti-symmetric deep
|
||||
graph network operator is called. (default: :obj:`1`)
|
||||
epsilon (float, optional): The discretization step size
|
||||
:math:`\epsilon`. (default: :obj:`0.1`)
|
||||
gamma (float, optional): The strength of the diffusion :math:`\gamma`.
|
||||
It regulates the stability of the method. (default: :obj:`0.1`)
|
||||
act (str, optional): The non-linear activation function :math:`\sigma`,
|
||||
*e.g.*, :obj:`"tanh"` or :obj:`"relu"`. (default: :class:`"tanh"`)
|
||||
act_kwargs (Dict[str, Any], optional): Arguments passed to the
|
||||
respective activation function defined by :obj:`act`.
|
||||
(default: :obj:`None`)
|
||||
bias (bool, optional): If set to :obj:`False`, the layer will not learn
|
||||
an additive bias. (default: :obj:`True`)
|
||||
|
||||
Shapes:
|
||||
- **input:**
|
||||
node features :math:`(|\mathcal{V}|, F_{in})`,
|
||||
edge indices :math:`(2, |\mathcal{E}|)`,
|
||||
edge weights :math:`(|\mathcal{E}|)` *(optional)*
|
||||
- **output:** node features :math:`(|\mathcal{V}|, F_{in})`
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
phi: Optional[MessagePassing] = None,
|
||||
num_iters: int = 1,
|
||||
epsilon: float = 0.1,
|
||||
gamma: float = 0.1,
|
||||
act: Union[str, Callable, None] = "tanh",
|
||||
act_kwargs: Optional[Dict[str, Any]] = None,
|
||||
bias: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.in_channels = in_channels
|
||||
self.num_iters = num_iters
|
||||
self.gamma = gamma
|
||||
self.epsilon = epsilon
|
||||
self.act = activation_resolver(act, **(act_kwargs or {}))
|
||||
|
||||
if phi is None:
|
||||
phi = GCNConv(in_channels, in_channels, bias=False)
|
||||
|
||||
self.W = Parameter(torch.empty(in_channels, in_channels))
|
||||
self.register_buffer("eye", torch.eye(in_channels))
|
||||
self.phi = phi
|
||||
|
||||
if bias:
|
||||
self.bias = Parameter(torch.empty(in_channels))
|
||||
else:
|
||||
self.register_parameter("bias", None)
|
||||
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self):
|
||||
r"""Resets all learnable parameters of the module."""
|
||||
torch.nn.init.kaiming_uniform_(self.W, a=math.sqrt(5))
|
||||
self.phi.reset_parameters()
|
||||
zeros(self.bias)
|
||||
|
||||
def forward(self, x: Tensor, edge_index: Adj, *args, **kwargs) -> Tensor:
|
||||
r"""Runs the forward pass of the module."""
|
||||
antisymmetric_W = self.W - self.W.t() - self.gamma * self.eye
|
||||
|
||||
for _ in range(self.num_iters):
|
||||
h = self.phi(x, edge_index, *args, **kwargs)
|
||||
h = x @ antisymmetric_W.t() + h
|
||||
|
||||
if self.bias is not None:
|
||||
h += self.bias
|
||||
|
||||
if self.act is not None:
|
||||
h = self.act(h)
|
||||
|
||||
x = x + self.epsilon * h
|
||||
|
||||
return x
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"{self.__class__.__name__}("
|
||||
f"{self.in_channels}, "
|
||||
f"phi={self.phi}, "
|
||||
f"num_iters={self.num_iters}, "
|
||||
f"epsilon={self.epsilon}, "
|
||||
f"gamma={self.gamma})"
|
||||
)
|
||||
|
||||
|
||||
model_cls = AntiSymmetricConv
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
node_features = torch.load("node_features.pt")
|
||||
edge_index = torch.load("edge_index.pt")
|
||||
|
||||
# Model instantiation and forward pass
|
||||
model = AntiSymmetricConv(in_channels=node_features.size(-1))
|
||||
output = model(node_features, edge_index)
|
||||
|
||||
# Save output to a file
|
||||
torch.save(output, "gt_output.pt")
|
||||
@@ -0,0 +1,89 @@
|
||||
import copy
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch_geometric.nn.conv import MessagePassing
|
||||
|
||||
|
||||
class DirGNNConv(torch.nn.Module):
|
||||
r"""A generic wrapper for computing graph convolution on directed
|
||||
graphs as described in the `"Edge Directionality Improves Learning on
|
||||
Heterophilic Graphs" <https://arxiv.org/abs/2305.10498>`_ paper.
|
||||
:class:`DirGNNConv` will pass messages both from source nodes to target
|
||||
nodes and from target nodes to source nodes.
|
||||
|
||||
Args:
|
||||
conv (MessagePassing): The underlying
|
||||
:class:`~torch_geometric.nn.conv.MessagePassing` layer to use.
|
||||
alpha (float, optional): The alpha coefficient used to weight the
|
||||
aggregations of in- and out-edges as part of a convex combination.
|
||||
(default: :obj:`0.5`)
|
||||
root_weight (bool, optional): If set to :obj:`True`, the layer will add
|
||||
transformed root node features to the output.
|
||||
(default: :obj:`True`)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
conv: MessagePassing,
|
||||
alpha: float = 0.5,
|
||||
root_weight: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.alpha = alpha
|
||||
self.root_weight = root_weight
|
||||
|
||||
self.conv_in = copy.deepcopy(conv)
|
||||
self.conv_out = copy.deepcopy(conv)
|
||||
|
||||
if hasattr(conv, "add_self_loops"):
|
||||
self.conv_in.add_self_loops = False
|
||||
self.conv_out.add_self_loops = False
|
||||
if hasattr(conv, "root_weight"):
|
||||
self.conv_in.root_weight = False
|
||||
self.conv_out.root_weight = False
|
||||
|
||||
if root_weight:
|
||||
self.lin = torch.nn.Linear(conv.in_channels, conv.out_channels)
|
||||
else:
|
||||
self.lin = None
|
||||
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self):
|
||||
r"""Resets all learnable parameters of the module."""
|
||||
self.conv_in.reset_parameters()
|
||||
self.conv_out.reset_parameters()
|
||||
if self.lin is not None:
|
||||
self.lin.reset_parameters()
|
||||
|
||||
def forward(self, x: Tensor, edge_index: Tensor) -> Tensor:
|
||||
"""""" # noqa: D419
|
||||
x_in = self.conv_in(x, edge_index)
|
||||
x_out = self.conv_out(x, edge_index.flip([0]))
|
||||
|
||||
out = self.alpha * x_out + (1 - self.alpha) * x_in
|
||||
|
||||
if self.root_weight:
|
||||
out = out + self.lin(x)
|
||||
|
||||
return out
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"{self.__class__.__name__}({self.conv_in}, alpha={self.alpha})"
|
||||
|
||||
|
||||
model_cls = DirGNNConv
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
node_features = torch.load("node_features.pt")
|
||||
edge_index = torch.load("edge_index.pt")
|
||||
|
||||
# Model instantiation and forward pass
|
||||
model = DirGNNConv(MessagePassing())
|
||||
output = model(node_features, edge_index)
|
||||
|
||||
# Save output to a file
|
||||
torch.save(output, "gt_output.pt")
|
||||
@@ -0,0 +1,198 @@
|
||||
import inspect
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
from torch.nn import Dropout, Linear, Sequential
|
||||
from torch_geometric.nn.attention import PerformerAttention
|
||||
from torch_geometric.nn.conv import MessagePassing
|
||||
from torch_geometric.nn.inits import reset
|
||||
from torch_geometric.nn.resolver import activation_resolver, normalization_resolver
|
||||
from torch_geometric.typing import Adj
|
||||
from torch_geometric.utils import to_dense_batch
|
||||
|
||||
|
||||
class GPSConv(torch.nn.Module):
|
||||
r"""The general, powerful, scalable (GPS) graph transformer layer from the
|
||||
`"Recipe for a General, Powerful, Scalable Graph Transformer"
|
||||
<https://arxiv.org/abs/2205.12454>`_ paper.
|
||||
|
||||
The GPS layer is based on a 3-part recipe:
|
||||
|
||||
1. Inclusion of positional (PE) and structural encodings (SE) to the input
|
||||
features (done in a pre-processing step via
|
||||
:class:`torch_geometric.transforms`).
|
||||
2. A local message passing layer (MPNN) that operates on the input graph.
|
||||
3. A global attention layer that operates on the entire graph.
|
||||
|
||||
.. note::
|
||||
|
||||
For an example of using :class:`GPSConv`, see
|
||||
`examples/graph_gps.py
|
||||
<https://github.com/pyg-team/pytorch_geometric/blob/master/examples/
|
||||
graph_gps.py>`_.
|
||||
|
||||
Args:
|
||||
channels (int): Size of each input sample.
|
||||
conv (MessagePassing, optional): The local message passing layer.
|
||||
heads (int, optional): Number of multi-head-attentions.
|
||||
(default: :obj:`1`)
|
||||
dropout (float, optional): Dropout probability of intermediate
|
||||
embeddings. (default: :obj:`0.`)
|
||||
act (str or Callable, optional): The non-linear activation function to
|
||||
use. (default: :obj:`"relu"`)
|
||||
act_kwargs (Dict[str, Any], optional): Arguments passed to the
|
||||
respective activation function defined by :obj:`act`.
|
||||
(default: :obj:`None`)
|
||||
norm (str or Callable, optional): The normalization function to
|
||||
use. (default: :obj:`"batch_norm"`)
|
||||
norm_kwargs (Dict[str, Any], optional): Arguments passed to the
|
||||
respective normalization function defined by :obj:`norm`.
|
||||
(default: :obj:`None`)
|
||||
attn_type (str): Global attention type, :obj:`multihead` or
|
||||
:obj:`performer`. (default: :obj:`multihead`)
|
||||
attn_kwargs (Dict[str, Any], optional): Arguments passed to the
|
||||
attention layer. (default: :obj:`None`)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
channels: int,
|
||||
conv: Optional[MessagePassing],
|
||||
heads: int = 1,
|
||||
dropout: float = 0.0,
|
||||
act: str = "relu",
|
||||
act_kwargs: Optional[Dict[str, Any]] = None,
|
||||
norm: Optional[str] = "batch_norm",
|
||||
norm_kwargs: Optional[Dict[str, Any]] = None,
|
||||
attn_type: str = "multihead",
|
||||
attn_kwargs: Optional[Dict[str, Any]] = None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.channels = channels
|
||||
self.conv = conv
|
||||
self.heads = heads
|
||||
self.dropout = dropout
|
||||
self.attn_type = attn_type
|
||||
|
||||
attn_kwargs = attn_kwargs or {}
|
||||
if attn_type == "multihead":
|
||||
self.attn = torch.nn.MultiheadAttention(
|
||||
channels,
|
||||
heads,
|
||||
batch_first=True,
|
||||
**attn_kwargs,
|
||||
)
|
||||
elif attn_type == "performer":
|
||||
self.attn = PerformerAttention(
|
||||
channels=channels,
|
||||
heads=heads,
|
||||
**attn_kwargs,
|
||||
)
|
||||
else:
|
||||
# TODO: Support BigBird
|
||||
raise ValueError(f"{attn_type} is not supported")
|
||||
|
||||
self.mlp = Sequential(
|
||||
Linear(channels, channels * 2),
|
||||
activation_resolver(act, **(act_kwargs or {})),
|
||||
Dropout(dropout),
|
||||
Linear(channels * 2, channels),
|
||||
Dropout(dropout),
|
||||
)
|
||||
|
||||
norm_kwargs = norm_kwargs or {}
|
||||
self.norm1 = normalization_resolver(norm, channels, **norm_kwargs)
|
||||
self.norm2 = normalization_resolver(norm, channels, **norm_kwargs)
|
||||
self.norm3 = normalization_resolver(norm, channels, **norm_kwargs)
|
||||
|
||||
self.norm_with_batch = False
|
||||
if self.norm1 is not None:
|
||||
signature = inspect.signature(self.norm1.forward)
|
||||
self.norm_with_batch = "batch" in signature.parameters
|
||||
|
||||
def reset_parameters(self):
|
||||
r"""Resets all learnable parameters of the module."""
|
||||
if self.conv is not None:
|
||||
self.conv.reset_parameters()
|
||||
self.attn._reset_parameters()
|
||||
reset(self.mlp)
|
||||
if self.norm1 is not None:
|
||||
self.norm1.reset_parameters()
|
||||
if self.norm2 is not None:
|
||||
self.norm2.reset_parameters()
|
||||
if self.norm3 is not None:
|
||||
self.norm3.reset_parameters()
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: Tensor,
|
||||
edge_index: Adj,
|
||||
batch: Optional[torch.Tensor] = None,
|
||||
**kwargs,
|
||||
) -> Tensor:
|
||||
r"""Runs the forward pass of the module."""
|
||||
hs = []
|
||||
if self.conv is not None: # Local MPNN.
|
||||
h = self.conv(x, edge_index, **kwargs)
|
||||
h = F.dropout(h, p=self.dropout, training=self.training)
|
||||
h = h + x
|
||||
if self.norm1 is not None:
|
||||
if self.norm_with_batch:
|
||||
h = self.norm1(h, batch=batch)
|
||||
else:
|
||||
h = self.norm1(h)
|
||||
hs.append(h)
|
||||
|
||||
# Global attention transformer-style model.
|
||||
h, mask = to_dense_batch(x, batch)
|
||||
|
||||
if isinstance(self.attn, torch.nn.MultiheadAttention):
|
||||
h, _ = self.attn(h, h, h, key_padding_mask=~mask, need_weights=False)
|
||||
elif isinstance(self.attn, PerformerAttention):
|
||||
h = self.attn(h, mask=mask)
|
||||
|
||||
h = h[mask]
|
||||
h = F.dropout(h, p=self.dropout, training=self.training)
|
||||
h = h + x # Residual connection.
|
||||
if self.norm2 is not None:
|
||||
if self.norm_with_batch:
|
||||
h = self.norm2(h, batch=batch)
|
||||
else:
|
||||
h = self.norm2(h)
|
||||
hs.append(h)
|
||||
|
||||
out = sum(hs) # Combine local and global outputs.
|
||||
|
||||
out = out + self.mlp(out)
|
||||
if self.norm3 is not None:
|
||||
if self.norm_with_batch:
|
||||
out = self.norm3(out, batch=batch)
|
||||
else:
|
||||
out = self.norm3(out)
|
||||
|
||||
return out
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"{self.__class__.__name__}({self.channels}, "
|
||||
f"conv={self.conv}, heads={self.heads}, "
|
||||
f"attn_type={self.attn_type})"
|
||||
)
|
||||
|
||||
|
||||
model_cls = GPSConv
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
node_features = torch.load("node_features.pt")
|
||||
edge_index = torch.load("edge_index.pt")
|
||||
|
||||
# Model instantiation and forward pass
|
||||
model = GPSConv(channels=node_features.size(-1), conv=MessagePassing())
|
||||
output = model(node_features, edge_index)
|
||||
|
||||
# Save output to a file
|
||||
torch.save(output, "gt_output.pt")
|
||||
@@ -0,0 +1,187 @@
|
||||
import math
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch.nn import BatchNorm1d, Parameter
|
||||
from torch_geometric.nn import inits
|
||||
from torch_geometric.nn.conv import MessagePassing
|
||||
from torch_geometric.nn.models import MLP
|
||||
from torch_geometric.typing import Adj, OptTensor
|
||||
from torch_geometric.utils import spmm
|
||||
|
||||
|
||||
class SparseLinear(MessagePassing):
|
||||
def __init__(self, in_channels: int, out_channels: int, bias: bool = True):
|
||||
super().__init__(aggr="add")
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
|
||||
self.weight = Parameter(torch.empty(in_channels, out_channels))
|
||||
if bias:
|
||||
self.bias = Parameter(torch.empty(out_channels))
|
||||
else:
|
||||
self.register_parameter("bias", None)
|
||||
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self):
|
||||
inits.kaiming_uniform(self.weight, fan=self.in_channels, a=math.sqrt(5))
|
||||
inits.uniform(self.in_channels, self.bias)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
edge_index: Adj,
|
||||
edge_weight: OptTensor = None,
|
||||
) -> Tensor:
|
||||
# propagate_type: (weight: Tensor, edge_weight: OptTensor)
|
||||
out = self.propagate(edge_index, weight=self.weight, edge_weight=edge_weight)
|
||||
|
||||
if self.bias is not None:
|
||||
out = out + self.bias
|
||||
|
||||
return out
|
||||
|
||||
def message(self, weight_j: Tensor, edge_weight: OptTensor) -> Tensor:
|
||||
if edge_weight is None:
|
||||
return weight_j
|
||||
else:
|
||||
return edge_weight.view(-1, 1) * weight_j
|
||||
|
||||
def message_and_aggregate(self, adj_t: Adj, weight: Tensor) -> Tensor:
|
||||
return spmm(adj_t, weight, reduce=self.aggr)
|
||||
|
||||
|
||||
class LINKX(torch.nn.Module):
|
||||
r"""The LINKX model from the `"Large Scale Learning on Non-Homophilous
|
||||
Graphs: New Benchmarks and Strong Simple Methods"
|
||||
<https://arxiv.org/abs/2110.14446>`_ paper.
|
||||
|
||||
.. math::
|
||||
\mathbf{H}_{\mathbf{A}} &= \textrm{MLP}_{\mathbf{A}}(\mathbf{A})
|
||||
|
||||
\mathbf{H}_{\mathbf{X}} &= \textrm{MLP}_{\mathbf{X}}(\mathbf{X})
|
||||
|
||||
\mathbf{Y} &= \textrm{MLP}_{f} \left( \sigma \left( \mathbf{W}
|
||||
[\mathbf{H}_{\mathbf{A}}, \mathbf{H}_{\mathbf{X}}] +
|
||||
\mathbf{H}_{\mathbf{A}} + \mathbf{H}_{\mathbf{X}} \right) \right)
|
||||
|
||||
.. note::
|
||||
|
||||
For an example of using LINKX, see `examples/linkx.py <https://
|
||||
github.com/pyg-team/pytorch_geometric/blob/master/examples/linkx.py>`_.
|
||||
|
||||
Args:
|
||||
num_nodes (int): The number of nodes in the graph.
|
||||
in_channels (int): Size of each input sample, or :obj:`-1` to derive
|
||||
the size from the first input(s) to the forward method.
|
||||
hidden_channels (int): Size of each hidden sample.
|
||||
out_channels (int): Size of each output sample.
|
||||
num_layers (int): Number of layers of :math:`\textrm{MLP}_{f}`.
|
||||
num_edge_layers (int, optional): Number of layers of
|
||||
:math:`\textrm{MLP}_{\mathbf{A}}`. (default: :obj:`1`)
|
||||
num_node_layers (int, optional): Number of layers of
|
||||
:math:`\textrm{MLP}_{\mathbf{X}}`. (default: :obj:`1`)
|
||||
dropout (float, optional): Dropout probability of each hidden
|
||||
embedding. (default: :obj:`0.0`)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_nodes: int,
|
||||
in_channels: int,
|
||||
hidden_channels: int,
|
||||
out_channels: int,
|
||||
num_layers: int,
|
||||
num_edge_layers: int = 1,
|
||||
num_node_layers: int = 1,
|
||||
dropout: float = 0.0,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.num_nodes = num_nodes
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
self.num_edge_layers = num_edge_layers
|
||||
|
||||
self.edge_lin = SparseLinear(num_nodes, hidden_channels)
|
||||
|
||||
if self.num_edge_layers > 1:
|
||||
self.edge_norm = BatchNorm1d(hidden_channels)
|
||||
channels = [hidden_channels] * num_edge_layers
|
||||
self.edge_mlp = MLP(channels, dropout=0.0, act_first=True)
|
||||
else:
|
||||
self.edge_norm = None
|
||||
self.edge_mlp = None
|
||||
|
||||
channels = [in_channels] + [hidden_channels] * num_node_layers
|
||||
self.node_mlp = MLP(channels, dropout=0.0, act_first=True)
|
||||
|
||||
self.cat_lin1 = torch.nn.Linear(hidden_channels, hidden_channels)
|
||||
self.cat_lin2 = torch.nn.Linear(hidden_channels, hidden_channels)
|
||||
|
||||
channels = [hidden_channels] * num_layers + [out_channels]
|
||||
self.final_mlp = MLP(channels, dropout=dropout, act_first=True)
|
||||
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self):
|
||||
r"""Resets all learnable parameters of the module."""
|
||||
self.edge_lin.reset_parameters()
|
||||
if self.edge_norm is not None:
|
||||
self.edge_norm.reset_parameters()
|
||||
if self.edge_mlp is not None:
|
||||
self.edge_mlp.reset_parameters()
|
||||
self.node_mlp.reset_parameters()
|
||||
self.cat_lin1.reset_parameters()
|
||||
self.cat_lin2.reset_parameters()
|
||||
self.final_mlp.reset_parameters()
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: OptTensor,
|
||||
edge_index: Adj,
|
||||
edge_weight: OptTensor = None,
|
||||
) -> Tensor:
|
||||
"""""" # noqa: D419
|
||||
out = self.edge_lin(edge_index, edge_weight)
|
||||
|
||||
if self.edge_norm is not None and self.edge_mlp is not None:
|
||||
out = out.relu_()
|
||||
out = self.edge_norm(out)
|
||||
out = self.edge_mlp(out)
|
||||
|
||||
out = out + self.cat_lin1(out)
|
||||
|
||||
if x is not None:
|
||||
x = self.node_mlp(x)
|
||||
out = out + x
|
||||
out = out + self.cat_lin2(x)
|
||||
|
||||
return self.final_mlp(out.relu_())
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"{self.__class__.__name__}(num_nodes={self.num_nodes}, "
|
||||
f"in_channels={self.in_channels}, "
|
||||
f"out_channels={self.out_channels})"
|
||||
)
|
||||
|
||||
|
||||
model_cls = LINKX
|
||||
|
||||
if __name__ == "__main__":
|
||||
node_features = torch.load("node_features.pt")
|
||||
edge_index = torch.load("edge_index.pt")
|
||||
|
||||
# Model instantiation and forward pass
|
||||
model = LINKX(
|
||||
num_nodes=node_features.size(0),
|
||||
in_channels=node_features.size(1),
|
||||
hidden_channels=node_features.size(1),
|
||||
out_channels=node_features.size(1),
|
||||
num_layers=1,
|
||||
)
|
||||
output = model(node_features, edge_index)
|
||||
|
||||
# Save output to a file
|
||||
torch.save(output, "gt_output.pt")
|
||||
@@ -0,0 +1,118 @@
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
from torch_geometric.nn import SimpleConv
|
||||
from torch_geometric.nn.dense.linear import Linear
|
||||
|
||||
|
||||
class PMLP(torch.nn.Module):
|
||||
r"""The P(ropagational)MLP model from the `"Graph Neural Networks are
|
||||
Inherently Good Generalizers: Insights by Bridging GNNs and MLPs"
|
||||
<https://arxiv.org/abs/2212.09034>`_ paper.
|
||||
:class:`PMLP` is identical to a standard MLP during training, but then
|
||||
adopts a GNN architecture during testing.
|
||||
|
||||
Args:
|
||||
in_channels (int): Size of each input sample.
|
||||
hidden_channels (int): Size of each hidden sample.
|
||||
out_channels (int): Size of each output sample.
|
||||
num_layers (int): The number of layers.
|
||||
dropout (float, optional): Dropout probability of each hidden
|
||||
embedding. (default: :obj:`0.`)
|
||||
norm (bool, optional): If set to :obj:`False`, will not apply batch
|
||||
normalization. (default: :obj:`True`)
|
||||
bias (bool, optional): If set to :obj:`False`, the module
|
||||
will not learn additive biases. (default: :obj:`True`)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
hidden_channels: int,
|
||||
out_channels: int,
|
||||
num_layers: int,
|
||||
dropout: float = 0.0,
|
||||
norm: bool = True,
|
||||
bias: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.in_channels = in_channels
|
||||
self.hidden_channels = hidden_channels
|
||||
self.out_channels = out_channels
|
||||
self.num_layers = num_layers
|
||||
self.dropout = dropout
|
||||
self.bias = bias
|
||||
|
||||
self.lins = torch.nn.ModuleList()
|
||||
self.lins.append(Linear(in_channels, hidden_channels, self.bias))
|
||||
for _ in range(self.num_layers - 2):
|
||||
lin = Linear(hidden_channels, hidden_channels, self.bias)
|
||||
self.lins.append(lin)
|
||||
self.lins.append(Linear(hidden_channels, out_channels, self.bias))
|
||||
|
||||
self.norm = None
|
||||
if norm:
|
||||
self.norm = torch.nn.BatchNorm1d(
|
||||
hidden_channels,
|
||||
affine=False,
|
||||
track_running_stats=False,
|
||||
)
|
||||
|
||||
self.conv = SimpleConv(aggr="mean", combine_root="self_loop")
|
||||
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self):
|
||||
r"""Resets all learnable parameters of the module."""
|
||||
for lin in self.lins:
|
||||
torch.nn.init.xavier_uniform_(lin.weight, gain=1.414)
|
||||
if self.bias:
|
||||
torch.nn.init.zeros_(lin.bias)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
edge_index: Optional[Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
"""""" # noqa: D419
|
||||
if not self.training and edge_index is None:
|
||||
raise ValueError(f"'edge_index' needs to be present during " f"inference in '{self.__class__.__name__}'")
|
||||
|
||||
for i in range(self.num_layers):
|
||||
x = x @ self.lins[i].weight.t()
|
||||
if not self.training:
|
||||
x = self.conv(x, edge_index)
|
||||
if self.bias:
|
||||
x = x + self.lins[i].bias
|
||||
if i != self.num_layers - 1:
|
||||
if self.norm is not None:
|
||||
x = self.norm(x)
|
||||
x = x.relu()
|
||||
x = F.dropout(x, p=self.dropout, training=self.training)
|
||||
|
||||
return x
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"{self.__class__.__name__}({self.in_channels}, " f"{self.out_channels}, num_layers={self.num_layers})"
|
||||
|
||||
|
||||
model_cls = PMLP
|
||||
|
||||
if __name__ == "__main__":
|
||||
node_features = torch.load("node_features.pt")
|
||||
edge_index = torch.load("edge_index.pt")
|
||||
|
||||
# Model instantiation and forward pass
|
||||
model = PMLP(
|
||||
in_channels=node_features.size(-1),
|
||||
hidden_channels=node_features.size(-1),
|
||||
out_channels=node_features.size(-1),
|
||||
num_layers=1,
|
||||
)
|
||||
output = model(node_features, edge_index)
|
||||
|
||||
# Save output to a file
|
||||
torch.save(output, "gt_output.pt")
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,13 @@
|
||||
from pathlib import Path
|
||||
|
||||
from pydantic_settings import BaseSettings
|
||||
|
||||
|
||||
class ModelImplSettings(BaseSettings):
|
||||
workspace_path: Path = Path("./git_ignore_folder/model_imp_workspace/") # Added type annotation for work_space
|
||||
|
||||
class Config:
|
||||
env_prefix = "MODEL_IMPL_" # Use MODEL_IMPL_ as prefix for environment variables
|
||||
|
||||
|
||||
MODEL_IMPL_SETTINGS = ModelImplSettings()
|
||||
@@ -0,0 +1,59 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
def shape_evaluator(target, prediction):
|
||||
if target is None or prediction is None:
|
||||
return None, 0
|
||||
tar_shape = target.shape
|
||||
pre_shape = prediction.shape
|
||||
|
||||
diff = []
|
||||
for i in range(max(len(tar_shape), len(pre_shape))):
|
||||
dim_tar = tar_shape[i] if i < len(tar_shape) else 0
|
||||
dim_pre = pre_shape[i] if i < len(pre_shape) else 0
|
||||
diff.append(abs(dim_tar - dim_pre))
|
||||
|
||||
metric = 1 / (np.exp(np.mean(diff)) + 1)
|
||||
return diff, metric
|
||||
|
||||
|
||||
def reshape_tensor(original_tensor, target_shape):
|
||||
new_tensor = torch.zeros(target_shape)
|
||||
for i, dim in enumerate(original_tensor.shape):
|
||||
new_tensor = new_tensor.narrow(i, 0, dim).copy_(original_tensor)
|
||||
|
||||
return new_tensor
|
||||
|
||||
|
||||
def value_evaluator(target, prediction):
|
||||
if target is None or prediction is None:
|
||||
return None, 0
|
||||
tar_shape = target.shape
|
||||
pre_shape = prediction.shape
|
||||
|
||||
# Determine the shape of the padded tensors
|
||||
dims = [
|
||||
max(s1, s2)
|
||||
for s1, s2 in zip(
|
||||
tar_shape + (1,) * (len(pre_shape) - len(tar_shape)),
|
||||
pre_shape + (1,) * (len(tar_shape) - len(pre_shape)),
|
||||
)
|
||||
]
|
||||
# Reshape both tensors to the determined shape
|
||||
target = target.reshape(*tar_shape, *(1,) * (max(len(tar_shape), len(pre_shape)) - len(tar_shape)))
|
||||
prediction = prediction.reshape(*pre_shape, *(1,) * (max(len(tar_shape), len(pre_shape)) - len(pre_shape)))
|
||||
target_padded = reshape_tensor(target, dims)
|
||||
prediction_padded = reshape_tensor(prediction, dims)
|
||||
|
||||
# Calculate the mean absolute difference
|
||||
diff = torch.abs(target_padded - prediction_padded)
|
||||
metric = 1 / (1 + np.exp(torch.mean(diff).item()))
|
||||
return diff, metric
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tar = torch.rand(4, 5, 5)
|
||||
pre = torch.rand(4, 1)
|
||||
print(shape_evaluator(tar, pre))
|
||||
print(value_evaluator(tar, pre)[1])
|
||||
@@ -0,0 +1,136 @@
|
||||
"""
|
||||
This is just an exmaple.
|
||||
It will be replaced wtih a list of ground truth tasks.
|
||||
"""
|
||||
|
||||
import math
|
||||
from typing import Any, Callable, Dict, Optional, Union
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch.nn import Parameter
|
||||
from torch_geometric.nn.conv import GCNConv, MessagePassing
|
||||
from torch_geometric.nn.inits import zeros
|
||||
from torch_geometric.nn.resolver import activation_resolver
|
||||
from torch_geometric.typing import Adj
|
||||
|
||||
|
||||
class AntiSymmetricConv(torch.nn.Module):
|
||||
r"""The anti-symmetric graph convolutional operator from the
|
||||
`"Anti-Symmetric DGN: a stable architecture for Deep Graph Networks"
|
||||
<https://openreview.net/forum?id=J3Y7cgZOOS>`_ paper.
|
||||
|
||||
.. math::
|
||||
\mathbf{x}^{\prime}_i = \mathbf{x}_i + \epsilon \cdot \sigma \left(
|
||||
(\mathbf{W}-\mathbf{W}^T-\gamma \mathbf{I}) \mathbf{x}_i +
|
||||
\Phi(\mathbf{X}, \mathcal{N}_i) + \mathbf{b}\right),
|
||||
|
||||
where :math:`\Phi(\mathbf{X}, \mathcal{N}_i)` denotes a
|
||||
:class:`~torch.nn.conv.MessagePassing` layer.
|
||||
|
||||
Args:
|
||||
in_channels (int): Size of each input sample.
|
||||
phi (MessagePassing, optional): The message passing module
|
||||
:math:`\Phi`. If set to :obj:`None`, will use a
|
||||
:class:`~torch_geometric.nn.conv.GCNConv` layer as default.
|
||||
(default: :obj:`None`)
|
||||
num_iters (int, optional): The number of times the anti-symmetric deep
|
||||
graph network operator is called. (default: :obj:`1`)
|
||||
epsilon (float, optional): The discretization step size
|
||||
:math:`\epsilon`. (default: :obj:`0.1`)
|
||||
gamma (float, optional): The strength of the diffusion :math:`\gamma`.
|
||||
It regulates the stability of the method. (default: :obj:`0.1`)
|
||||
act (str, optional): The non-linear activation function :math:`\sigma`,
|
||||
*e.g.*, :obj:`"tanh"` or :obj:`"relu"`. (default: :class:`"tanh"`)
|
||||
act_kwargs (Dict[str, Any], optional): Arguments passed to the
|
||||
respective activation function defined by :obj:`act`.
|
||||
(default: :obj:`None`)
|
||||
bias (bool, optional): If set to :obj:`False`, the layer will not learn
|
||||
an additive bias. (default: :obj:`True`)
|
||||
|
||||
Shapes:
|
||||
- **input:**
|
||||
node features :math:`(|\mathcal{V}|, F_{in})`,
|
||||
edge indices :math:`(2, |\mathcal{E}|)`,
|
||||
edge weights :math:`(|\mathcal{E}|)` *(optional)*
|
||||
- **output:** node features :math:`(|\mathcal{V}|, F_{in})`
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
phi: Optional[MessagePassing] = None,
|
||||
num_iters: int = 1,
|
||||
epsilon: float = 0.1,
|
||||
gamma: float = 0.1,
|
||||
act: Union[str, Callable, None] = "tanh",
|
||||
act_kwargs: Optional[Dict[str, Any]] = None,
|
||||
bias: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.in_channels = in_channels
|
||||
self.num_iters = num_iters
|
||||
self.gamma = gamma
|
||||
self.epsilon = epsilon
|
||||
self.act = activation_resolver(act, **(act_kwargs or {}))
|
||||
|
||||
if phi is None:
|
||||
phi = GCNConv(in_channels, in_channels, bias=False)
|
||||
|
||||
self.W = Parameter(torch.empty(in_channels, in_channels))
|
||||
self.register_buffer("eye", torch.eye(in_channels))
|
||||
self.phi = phi
|
||||
|
||||
if bias:
|
||||
self.bias = Parameter(torch.empty(in_channels))
|
||||
else:
|
||||
self.register_parameter("bias", None)
|
||||
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self):
|
||||
r"""Resets all learnable parameters of the module."""
|
||||
torch.nn.init.kaiming_uniform_(self.W, a=math.sqrt(5))
|
||||
self.phi.reset_parameters()
|
||||
zeros(self.bias)
|
||||
|
||||
def forward(self, x: Tensor, edge_index: Adj, *args, **kwargs) -> Tensor:
|
||||
r"""Runs the forward pass of the module."""
|
||||
antisymmetric_W = self.W - self.W.t() - self.gamma * self.eye
|
||||
|
||||
for _ in range(self.num_iters):
|
||||
h = self.phi(x, edge_index, *args, **kwargs)
|
||||
h = x @ antisymmetric_W.t() + h
|
||||
|
||||
if self.bias is not None:
|
||||
h += self.bias
|
||||
|
||||
if self.act is not None:
|
||||
h = self.act(h)
|
||||
|
||||
x = x + self.epsilon * h
|
||||
|
||||
return x
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"{self.__class__.__name__}("
|
||||
f"{self.in_channels}, "
|
||||
f"phi={self.phi}, "
|
||||
f"num_iters={self.num_iters}, "
|
||||
f"epsilon={self.epsilon}, "
|
||||
f"gamma={self.gamma})"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
node_features = torch.load("node_features.pt")
|
||||
edge_index = torch.load("edge_index.pt")
|
||||
|
||||
# Model instantiation and forward pass
|
||||
model = AntiSymmetricConv(in_channels=node_features.size(-1))
|
||||
output = model(node_features, edge_index)
|
||||
|
||||
# Save output to a file
|
||||
torch.save(output, "gt_output.pt")
|
||||
@@ -0,0 +1,90 @@
|
||||
"""
|
||||
This file will be removed in the future and replaced by
|
||||
- rdagent/app/model_implementation/eval.py
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
# randomly generate a input graph, node_feature and edge_index
|
||||
# 1000 nodes, 128 dim node feature, 2000 edges
|
||||
import torch
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from rdagent.oai.llm_utils import APIBackend
|
||||
|
||||
assert load_dotenv()
|
||||
formula_info = {
|
||||
"name": "Anti-Symmetric Deep Graph Network (A-DGN)",
|
||||
"description": "A framework for stable and non-dissipative DGN design. It ensures long-range information preservation between nodes and prevents gradient vanishing or explosion during training.",
|
||||
"formulation": "x_u^{(l)} = x_u^{(l-1)} + \\epsilon \\sigma \\left( W^T x_u^{(l-1)} + \\Phi(X^{(l-1)}, N_u) + b \\right)",
|
||||
"variables": {
|
||||
"x_u^{(l)}": "The state of node u at layer l",
|
||||
"\\epsilon": "The step size in the Euler discretization",
|
||||
"\\sigma": "A monotonically non-decreasing activation function",
|
||||
"W": "An anti-symmetric weight matrix",
|
||||
"X^{(l-1)}": "The node feature matrix at layer l-1",
|
||||
"N_u": "The set of neighbors of node u",
|
||||
"b": "A bias vector",
|
||||
},
|
||||
}
|
||||
|
||||
system_prompt = "You are an assistant whose job is to answer user's question."
|
||||
user_prompt = "With the following given information, write a python code using pytorch and torch_geometric to implement the model. This model is in the graph learning field, only have one layer. The input will be node_feature [num_nodes, dim_feature] and edge_index [2, num_edges], and they should be loaded from the files 'node_features.pt' and 'edge_index.pt'. There is not edge attribute or edge weight as input. The model should detect the node_feature and edge_index shape, if there is Linear transformation layer in the model, the input and output shape should be consistent. The in_channels is the dimension of the node features. You code should contain additional 'if __name__ == '__main__', where you should load the node_feature and edge_index from the files and run the model, and save the output to a file 'llm_output.pt'. Implement the model forward function based on the following information: model formula information. 1. model name: {}, 2. model description: {}, 3. model formulation: {}, 4. model variables: {}. You must complete the forward function as far as you can do.".format(
|
||||
formula_info["name"],
|
||||
formula_info["description"],
|
||||
formula_info["formulation"],
|
||||
formula_info["variables"],
|
||||
)
|
||||
|
||||
resp = APIBackend(use_chat_cache=False).build_messages_and_create_chat_completion(user_prompt, system_prompt)
|
||||
|
||||
print(resp)
|
||||
|
||||
# take the code part from the response and save it to a file, the code is covered in the ```python``` block
|
||||
code = resp.split("```python")[1].split("```")[0]
|
||||
with open("llm_code.py", "w") as f:
|
||||
f.write(code)
|
||||
|
||||
average_shape_eval = []
|
||||
average_value_eval = []
|
||||
for test_mode in ["zeros", "ones", "randn"]:
|
||||
if test_mode == "zeros":
|
||||
node_feature = torch.zeros(1000, 128)
|
||||
elif test_mode == "ones":
|
||||
node_feature = torch.ones(1000, 128)
|
||||
elif test_mode == "randn":
|
||||
node_feature = torch.randn(1000, 128)
|
||||
edge_index = torch.randint(0, 1000, (2, 2000))
|
||||
|
||||
torch.save(node_feature, "node_features.pt")
|
||||
torch.save(edge_index, "edge_index.pt")
|
||||
|
||||
try:
|
||||
os.system("python llm_code.py")
|
||||
except:
|
||||
print("Error in running the LLM code")
|
||||
os.system("python gt_code.py")
|
||||
os.system("rm edge_index.pt")
|
||||
os.system("rm node_features.pt")
|
||||
# load the output and print the shape
|
||||
|
||||
from evaluator import shape_evaluator, value_evaluator
|
||||
|
||||
try:
|
||||
llm_output = torch.load("llm_output.pt")
|
||||
except:
|
||||
llm_output = None
|
||||
gt_output = torch.load("gt_output.pt")
|
||||
|
||||
average_shape_eval.append(shape_evaluator(llm_output, gt_output)[1])
|
||||
average_value_eval.append(value_evaluator(llm_output, gt_output)[1])
|
||||
|
||||
print("Shape evaluation: ", average_shape_eval[-1])
|
||||
print("Value evaluation:super().generate(task_l) ", average_value_eval[-1])
|
||||
|
||||
os.system("rm llm_output.pt")
|
||||
os.system("rm gt_output.pt")
|
||||
os.system("rm llm_code.py")
|
||||
|
||||
print("Average shape evaluation: ", sum(average_shape_eval) / len(average_shape_eval))
|
||||
print("Average value evaluation: ", sum(average_value_eval) / len(average_value_eval))
|
||||
@@ -0,0 +1,210 @@
|
||||
import json
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Dict, Optional, Sequence
|
||||
|
||||
import torch
|
||||
|
||||
from rdagent.components.coder.model_coder.conf import MODEL_IMPL_SETTINGS
|
||||
from rdagent.components.loader.task_loader import ModelTaskLoader
|
||||
from rdagent.core.exception import CodeFormatException
|
||||
from rdagent.core.experiment import Experiment, FBImplementation, ImpLoader, Task
|
||||
from rdagent.utils import get_module_by_module_path
|
||||
|
||||
|
||||
class ModelTask(Task):
|
||||
# TODO: it should change when the Task changes.
|
||||
name: str
|
||||
description: str
|
||||
formulation: str
|
||||
variables: Dict[str, str] # map the variable name to the variable description
|
||||
|
||||
def __init__(
|
||||
self, name: str, description: str, formulation: str, variables: Dict[str, str], key: Optional[str] = None
|
||||
) -> None:
|
||||
"""
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
key : Optional[str]
|
||||
Key is a string to identify the task.
|
||||
It will be used to connect to other information(e.g. ground truth).
|
||||
"""
|
||||
self.name = name
|
||||
self.description = description
|
||||
self.formulation = formulation
|
||||
self.variables = variables
|
||||
self.key = key
|
||||
|
||||
def get_information(self):
|
||||
return f"""name: {self.name}
|
||||
description: {self.description}
|
||||
formulation: {self.formulation}
|
||||
variables: {self.variables}
|
||||
key: {self.key}
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def from_dict(dict):
|
||||
return ModelTask(**dict)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<{self.__class__.__name__} {self.name}>"
|
||||
|
||||
|
||||
class ModelImplementation(FBImplementation):
|
||||
"""
|
||||
It is a Pytorch model implementation task;
|
||||
All the things are placed in a folder.
|
||||
|
||||
Folder
|
||||
- data source and documents prepared by `prepare`
|
||||
- Please note that new data may be passed in dynamically in `execute`
|
||||
- code (file `model.py` ) injected by `inject_code`
|
||||
- the `model.py` that contains a variable named `model_cls` which indicates the implemented model structure
|
||||
- `model_cls` is a instance of `torch.nn.Module`;
|
||||
|
||||
|
||||
We'll import the model in the implementation in file `model.py` after setting the cwd into the directory
|
||||
- from model import model_cls
|
||||
- initialize the model by initializing it `model_cls(input_dim=INPUT_DIM)`
|
||||
- And then verify the model.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, target_task: Task) -> None:
|
||||
super().__init__(target_task)
|
||||
self.path = None
|
||||
|
||||
def prepare(self) -> None:
|
||||
"""
|
||||
Prepare for the workspace;
|
||||
"""
|
||||
unique_id = uuid.uuid4()
|
||||
self.path = MODEL_IMPL_SETTINGS.workspace_path / f"M{unique_id}"
|
||||
# start with `M` so that it can be imported via python
|
||||
self.path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def execute(self, data=None, config: dict = {}):
|
||||
mod = get_module_by_module_path(str(self.path / "model.py"))
|
||||
try:
|
||||
model_cls = mod.model_cls
|
||||
except AttributeError:
|
||||
raise CodeFormatException("The model_cls is not implemented in the model.py")
|
||||
# model_init =
|
||||
|
||||
assert isinstance(data, tuple)
|
||||
node_feature, _ = data
|
||||
in_channels = node_feature.size(-1)
|
||||
m = model_cls(in_channels)
|
||||
|
||||
# TODO: initialize all the parameters of `m` to `model_eval_param_init`
|
||||
model_eval_param_init: float = config["model_eval_param_init"]
|
||||
|
||||
# initialize all parameters of `m` to `model_eval_param_init`
|
||||
for _, param in m.named_parameters():
|
||||
param.data.fill_(model_eval_param_init)
|
||||
|
||||
assert isinstance(data, tuple)
|
||||
return m(*data)
|
||||
|
||||
def execute_desc(self) -> str:
|
||||
return """
|
||||
The the implemented code will be placed in a file like <uuid>/model.py
|
||||
|
||||
We'll import the model in the implementation in file `model.py` after setting the cwd into the directory
|
||||
- from model import model_cls (So you must have a variable named `model_cls` in the file)
|
||||
- So your implemented code could follow the following pattern
|
||||
```Python
|
||||
class XXXLayer(torch.nn.Module):
|
||||
...
|
||||
model_cls = XXXLayer
|
||||
```
|
||||
- initialize the model by initializing it `model_cls(input_dim=INPUT_DIM)`
|
||||
- And then verify the model by comparing the output tensors by feeding specific input tensor.
|
||||
"""
|
||||
|
||||
|
||||
class ModelExperiment(Experiment[ModelTask, ModelImplementation]):
|
||||
...
|
||||
|
||||
|
||||
class ModelTaskLoaderJson(ModelTaskLoader):
|
||||
# def __init__(self, json_uri: str, select_model: Optional[str] = None) -> None:
|
||||
# super().__init__()
|
||||
# self.json_uri = json_uri
|
||||
# self.select_model = 'A-DGN'
|
||||
|
||||
# def load(self, *argT, **kwargs) -> Sequence[ModelImplTask]:
|
||||
# # json is supposed to be in the format of {model_name: dict{model_data}}
|
||||
# model_dict = json.load(open(self.json_uri, "r"))
|
||||
# if self.select_model is not None:
|
||||
# assert self.select_model in model_dict
|
||||
# model_name = self.select_model
|
||||
# model_data = model_dict[self.select_model]
|
||||
# else:
|
||||
# model_name, model_data = list(model_dict.items())[0]
|
||||
|
||||
# model_impl_task = ModelImplTask(
|
||||
# name=model_name,
|
||||
# description=model_data["description"],
|
||||
# formulation=model_data["formulation"],
|
||||
# variables=model_data["variables"],
|
||||
# key=model_name
|
||||
# )
|
||||
|
||||
# return [model_impl_task]
|
||||
|
||||
def __init__(self, json_uri: str) -> None:
|
||||
super().__init__()
|
||||
self.json_uri = json_uri
|
||||
|
||||
def load(self, *argT, **kwargs) -> Sequence[ModelTask]:
|
||||
# json is supposed to be in the format of {model_name: dict{model_data}}
|
||||
model_dict = json.load(open(self.json_uri, "r"))
|
||||
|
||||
# FIXME: the model in the json file is not right due to extraction error
|
||||
# We should fix them case by case in the future
|
||||
#
|
||||
# formula_info = {
|
||||
# "name": "Anti-Symmetric Deep Graph Network (A-DGN)",
|
||||
# "description": "A framework for stable and non-dissipative DGN design. It ensures long-range information preservation between nodes and prevents gradient vanishing or explosion during training.",
|
||||
# "formulation": r"\mathbf{x}^{\prime}_i = \mathbf{x}_i + \epsilon \cdot \sigma \left( (\mathbf{W}-\mathbf{W}^T-\gamma \mathbf{I}) \mathbf{x}_i + \Phi(\mathbf{X}, \mathcal{N}_i) + \mathbf{b}\right),",
|
||||
# "variables": {
|
||||
# r"\mathbf{x}_i": "The state of node i at previous layer",
|
||||
# r"\epsilon": "The step size in the Euler discretization",
|
||||
# r"\sigma": "A monotonically non-decreasing activation function",
|
||||
# r"\Phi": "A graph convolutional operator",
|
||||
# r"W": "An anti-symmetric weight matrix",
|
||||
# r"\mathbf{x}^{\prime}_i": "The node feature matrix at layer l-1",
|
||||
# r"\mathcal{N}_i": "The set of neighbors of node u",
|
||||
# r"\mathbf{b}": "A bias vector",
|
||||
# },
|
||||
# "key": "A-DGN",
|
||||
# }
|
||||
model_impl_task_list = []
|
||||
for model_name, model_data in model_dict.items():
|
||||
model_impl_task = ModelTask(
|
||||
name=model_name,
|
||||
description=model_data["description"],
|
||||
formulation=model_data["formulation"],
|
||||
variables=model_data["variables"],
|
||||
key=model_data["key"],
|
||||
)
|
||||
model_impl_task_list.append(model_impl_task)
|
||||
return model_impl_task_list
|
||||
|
||||
|
||||
class ModelImpLoader(ImpLoader[ModelTask, ModelImplementation]):
|
||||
def __init__(self, path: Path) -> None:
|
||||
self.path = Path(path)
|
||||
|
||||
def load(self, task: ModelTask) -> ModelImplementation:
|
||||
assert task.key is not None
|
||||
mti = ModelImplementation(task)
|
||||
mti.prepare()
|
||||
with open(self.path / f"{task.key}.py", "r") as f:
|
||||
code = f.read()
|
||||
mti.inject_code(**{"model.py": code})
|
||||
return mti
|
||||
@@ -0,0 +1,45 @@
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
from jinja2 import Environment, StrictUndefined
|
||||
|
||||
from rdagent.components.coder.model_coder.model import (
|
||||
ModelExperiment,
|
||||
ModelImplementation,
|
||||
)
|
||||
from rdagent.core.prompts import Prompts
|
||||
from rdagent.core.task_generator import TaskGenerator
|
||||
from rdagent.oai.llm_utils import APIBackend
|
||||
|
||||
DIRNAME = Path(__file__).absolute().resolve().parent
|
||||
|
||||
|
||||
class ModelCodeWriter(TaskGenerator[ModelExperiment]):
|
||||
def generate(self, exp: ModelExperiment) -> ModelExperiment:
|
||||
mti_l = []
|
||||
for t in exp.sub_tasks:
|
||||
mti = ModelImplementation(t)
|
||||
mti.prepare()
|
||||
pr = Prompts(file_path=DIRNAME / "prompt.yaml")
|
||||
|
||||
user_prompt_tpl = Environment(undefined=StrictUndefined).from_string(pr["code_implement_user"])
|
||||
sys_prompt_tpl = Environment(undefined=StrictUndefined).from_string(pr["code_implement_sys"])
|
||||
|
||||
user_prompt = user_prompt_tpl.render(
|
||||
name=t.name,
|
||||
description=t.description,
|
||||
formulation=t.formulation,
|
||||
variables=t.variables,
|
||||
execute_desc=mti.execute_desc(),
|
||||
)
|
||||
system_prompt = sys_prompt_tpl.render()
|
||||
|
||||
resp = APIBackend().build_messages_and_create_chat_completion(user_prompt, system_prompt)
|
||||
|
||||
# Extract the code part from the response
|
||||
match = re.search(r".*```[Pp]ython\n(.*)\n```.*", resp, re.DOTALL)
|
||||
code = match.group(1)
|
||||
mti.inject_code(**{"model.py": code})
|
||||
mti_l.append(mti)
|
||||
exp.sub_implementations = mti_l
|
||||
return exp
|
||||
@@ -0,0 +1,18 @@
|
||||
|
||||
|
||||
code_implement_sys: -|
|
||||
You are an assistant whose job is to answer user's question."
|
||||
code_implement_user: -|
|
||||
With the following given information, write a python code using pytorch and torch_geometric to implement the model.
|
||||
This model is in the graph learning field, only have one layer.
|
||||
The input will be node_feature [num_nodes, dim_feature] and edge_index [2, num_edges] (It would be the input of the forward model)
|
||||
There is not edge attribute or edge weight as input. The model should detect the node_feature and edge_index shape, if there is Linear transformation layer in the model, the input and output shape should be consistent. The in_channels is the dimension of the node features.
|
||||
Implement the model forward function based on the following information:model formula information.
|
||||
1. model name:{{name}}
|
||||
2. model description:{{description}}
|
||||
3. model formulation:{{formulation}}
|
||||
4. model variables:{{variables}}.
|
||||
You must complete the forward function as far as you can do.
|
||||
\# Execution
|
||||
Your implemented code will be exectued in the follow way
|
||||
{{execute_desc}}
|
||||
@@ -0,0 +1,8 @@
|
||||
extract_model_formulation_system: |-
|
||||
offer description of the proposed model in this paper, write a latex formula with variable of the model. the format should be like " "Model Name": {
|
||||
"description": "",
|
||||
"formulation": "",
|
||||
"variables": {
|
||||
"\\hat{y}_u": "The predicted output for node u",
|
||||
}"
|
||||
such format content should be begin with ```json and end with ``` and the content should be in json format.
|
||||
@@ -0,0 +1,134 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
from rdagent.components.coder.model_coder.model import ModelExperiment, ModelTask
|
||||
from rdagent.components.document_reader.document_reader import (
|
||||
load_and_process_pdfs_by_langchain,
|
||||
)
|
||||
from rdagent.components.loader.task_loader import ModelTaskLoader
|
||||
from rdagent.core.log import RDAgentLog
|
||||
from rdagent.core.prompts import Prompts
|
||||
from rdagent.oai.llm_utils import APIBackend
|
||||
|
||||
document_process_prompts = Prompts(file_path=Path(__file__).parent / "prompts.yaml")
|
||||
|
||||
|
||||
def extract_model_from_doc(doc_content: str) -> dict:
|
||||
"""
|
||||
Extract model information from document content.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
doc_content : str
|
||||
Document content.
|
||||
|
||||
Returns
|
||||
-------
|
||||
dict
|
||||
{factor_name: dict{description, formulation, variables}}
|
||||
"""
|
||||
session = APIBackend().build_chat_session(
|
||||
session_system_prompt=document_process_prompts["extract_model_formulation_system"],
|
||||
)
|
||||
current_user_prompt = doc_content
|
||||
|
||||
# Extract model information from document content.
|
||||
model_dict = {}
|
||||
|
||||
for _ in range(10):
|
||||
# try to extract model information from the document content, retry at most 10 times.
|
||||
extract_result_resp = session.build_chat_completion(
|
||||
user_prompt=current_user_prompt,
|
||||
json_mode=False,
|
||||
)
|
||||
re_search_res = re.search(r"```json(.*)```", extract_result_resp, re.S)
|
||||
ret_json_str = re_search_res.group(1) if re_search_res is not None else ""
|
||||
try:
|
||||
ret_dict = json.loads(ret_json_str)
|
||||
parse_success = bool(isinstance(ret_dict, dict))
|
||||
except json.JSONDecodeError:
|
||||
parse_success = False
|
||||
if ret_json_str is None or not parse_success:
|
||||
current_user_prompt = "Your response didn't follow the instruction might be wrong json format. Try again."
|
||||
else:
|
||||
for name, formulation_and_description in ret_dict.items():
|
||||
if name not in model_dict:
|
||||
model_dict[name] = formulation_and_description
|
||||
if len(model_dict) == 0:
|
||||
current_user_prompt = "No model extracted. Please try again."
|
||||
else:
|
||||
break
|
||||
|
||||
RDAgentLog().info(f"已经完成{len(model_dict)}个模型的提取")
|
||||
|
||||
return model_dict
|
||||
|
||||
|
||||
def merge_file_to_model_dict_to_model_dict(
|
||||
file_to_model_dict: dict[str, dict],
|
||||
) -> dict:
|
||||
model_dict = {}
|
||||
for file_name in file_to_model_dict:
|
||||
for model_name in file_to_model_dict[file_name]:
|
||||
model_dict.setdefault(model_name, [])
|
||||
model_dict[model_name].append(file_to_model_dict[file_name][model_name])
|
||||
|
||||
model_dict_simple_deduplication = {}
|
||||
for model_name in model_dict:
|
||||
if len(model_dict[model_name]) > 1:
|
||||
model_dict_simple_deduplication[model_name] = max(
|
||||
model_dict[model_name],
|
||||
key=lambda x: len(x["formulation"]),
|
||||
)
|
||||
else:
|
||||
model_dict_simple_deduplication[model_name] = model_dict[model_name][0]
|
||||
return model_dict_simple_deduplication
|
||||
|
||||
|
||||
def extract_model_from_docs(docs_dict):
|
||||
model_dict = {}
|
||||
for doc_name, doc_content in docs_dict.items():
|
||||
model_dict[doc_name] = extract_model_from_doc(doc_content)
|
||||
return model_dict
|
||||
|
||||
|
||||
class ModelExperimentLoaderFromDict(ModelTaskLoader):
|
||||
def load(self, model_dict: dict) -> list:
|
||||
"""Load data from a dict."""
|
||||
task_l = []
|
||||
for model_name, model_data in model_dict.items():
|
||||
task = ModelTask(
|
||||
name=model_name,
|
||||
description=model_data["description"],
|
||||
formulation=model_data["formulation"],
|
||||
variables=model_data["variables"],
|
||||
key=model_name,
|
||||
)
|
||||
task_l.append(task)
|
||||
return ModelExperiment(sub_tasks=task_l)
|
||||
|
||||
|
||||
class ModelExperimentLoaderFromPDFfiles(ModelTaskLoader):
|
||||
def load(self, file_or_folder_path: Path) -> dict:
|
||||
docs_dict = load_and_process_pdfs_by_langchain(Path(file_or_folder_path)) # dict{file_path:content}
|
||||
model_dict = extract_model_from_docs(
|
||||
docs_dict
|
||||
) # dict{file_name: dict{model_name: dict{description, formulation, variables}}}
|
||||
model_dict = merge_file_to_model_dict_to_model_dict(
|
||||
model_dict
|
||||
) # dict {model_name: dict{description, formulation, variables}}
|
||||
return ModelExperimentLoaderFromDict().load(model_dict)
|
||||
|
||||
|
||||
def main(path="../test_doc"):
|
||||
doc_dict = load_and_process_pdfs_by_langchain(Path(path))
|
||||
print(doc_dict.keys()) # if you run code like "python -u", the print content will be truncated
|
||||
|
||||
|
||||
import fire
|
||||
|
||||
if __name__ == "__main__":
|
||||
fire.Fire(main)
|
||||
Reference in New Issue
Block a user