diff --git a/autogl/module/nas/estimator/one_shot.py b/autogl/module/nas/estimator/one_shot.py index e887138..e43695c 100644 --- a/autogl/module/nas/estimator/one_shot.py +++ b/autogl/module/nas/estimator/one_shot.py @@ -3,7 +3,7 @@ import torch.nn.functional as F from ..space import BaseSpace from .base import BaseEstimator - +import torch class OneShotEstimator(BaseEstimator): """ @@ -18,4 +18,5 @@ class OneShotEstimator(BaseEstimator): pred = model(dset)[getattr(dset, f"{mask}_mask")] y = dset.y[getattr(dset, f'{mask}_mask')] loss = F.nll_loss(pred, y) - return loss, loss + acc=sum(pred.max(1)[1]==y).item()/y.size(0) + return acc, loss diff --git a/autogl/module/nas/space/graph_nas.py b/autogl/module/nas/space/graph_nas.py index 95c290f..8e392d0 100644 --- a/autogl/module/nas/space/graph_nas.py +++ b/autogl/module/nas/space/graph_nas.py @@ -114,7 +114,11 @@ def gnn_map(gnn_name, in_dim, out_dim, concat=False, bias=True) -> nn.Module: elif gnn_name == "linear": return LinearConv(in_dim, out_dim, bias=bias) elif gnn_name == "zero": - return ZeroConv(in_dim, out_dim, bias=bias) + # return ZeroConv(in_dim, out_dim, bias=bias) + return Identity() +class Identity(nn.Module): + def forward(self, x, edge_index, edge_weight=None): + return x class LinearConv(nn.Module): def __init__(self, in_channels, @@ -207,6 +211,8 @@ class GraphNasNodeClassificationSpace(BaseSpace): setattr(self,f"act",self.setLayerChoice(2*layer,[act_map_nn(a)for a in act_list],key=f"act")) setattr(self,f"concat",self.setLayerChoice(2*layer+1,map_nn(["add", "product", "concat"]) ,key=f"concat")) self._initialized = True + self.classifier1 = nn.Linear(self.hidden_dim*self.layer_number, self.output_dim) + self.classifier2 = nn.Linear(self.hidden_dim, self.output_dim) def forward(self, data): x, edges = data.x, data.edge_index # x [2708,1433] ,[2, 10556] @@ -220,6 +226,7 @@ class GraphNasNodeClassificationSpace(BaseSpace): x = torch.cat(prev_nodes_out[2:],dim=1) x = F.leaky_relu(x) x = F.dropout(x, p=self.dropout, training = self.training) + x = self.classifier1(x) else: act=getattr(self, f"act") con=getattr(self, f"concat")() @@ -236,6 +243,11 @@ class GraphNasNodeClassificationSpace(BaseSpace): x=tmp x = act(x) x = F.dropout(x, p=self.dropout, training = self.training) + if con=='concat': + x=self.classifier1(x) + else: + x=self.classifier2(x) + # set_trace() return F.log_softmax(x, dim=1) def export(self, selection, device) -> BaseModel: