modify code of the file mindspore_hub_conf.py which network is gat and mass

model_zoo/official/gnn/gat/mindspore_hub_conf.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Huawei Technologies Co., Ltd
+# Copyright 2020-2021 Huawei Technologies Co., Ltd
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -14,12 +14,34 @@
 # ============================================================================
 """hub config."""
 from src.gat import GAT
-
+from src.config import GatConfig
-def gat(*args, **kwargs):
-    return GAT(*args, **kwargs)
-
 
 def create_network(name, *args, **kwargs):
+    """ create net work"""
     if name == "gat":
-        return gat(*args, **kwargs)
+
+        if "ftr_dims" in kwargs:
+            featureDims = kwargs.get("ftr_dims")
+        else:
+            featureDims = 3706
+
+        if "num_class" in kwargs:
+            numClass = kwargs.get("num_class")
+        else:
+            numClass = 10
+
+        if "num_nodes" in kwargs:
+            numNodes = kwargs.get("num_nodes")
+        else:
+            numNodes = 30
+
+        gat_net = GAT(featureDims,
+                      numClass,
+                      numNodes,
+                      GatConfig.hid_units,
+                      GatConfig.n_heads,
+                      attn_drop=GatConfig.attn_dropout,
+                      ftr_drop=GatConfig.feature_dropout)
+
+        return gat_net
     raise NotImplementedError(f"{name} is not implemented in the repo")

model_zoo/official/nlp/mass/mindspore_hub_conf.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Huawei Technologies Co., Ltd
+# Copyright 2020-2021 Huawei Technologies Co., Ltd
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -13,8 +13,8 @@
 # limitations under the License.
 # ============================================================================
 """hub config."""
+import os
 import mindspore.common.dtype as mstype
-
 from config import TransformerConfig
 from src.transformer import TransformerNetworkWithLoss, TransformerInferModel
 
@@ -27,10 +27,10 @@ def get_config(config):
 def create_network(name, *args, **kwargs):
     """create mass network."""
     if name == "mass":
-        if "config" in kwargs:
+        # get the config running dir
-            config = get_config(kwargs["config"])
+        configDir = os.path.split(os.path.realpath(__file__))[0] + "/config/config.json"
-        else:
+        # get the config
-            raise NotImplementedError(f"Please make sure the configuration file path is correct")
+        config = get_config(configDir)
         is_training = kwargs.get("is_training", False)
         if is_training:
             return TransformerNetworkWithLoss(config, is_training=is_training, *args)