fix classifier; fix zero conv

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

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: