Upload model

config.json

@@ -0,0 +1,17 @@
+{
+  "architectures": [
+    "GCNModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_gcn.GCNConfig",
+    "AutoModel": "modeling_gcn.GCNModel"
+  },
+  "emb_input": 20,
+  "hidden_size": 64,
+  "input_feature": 64,
+  "model_type": "gcn",
+  "n_layers": 6,
+  "num_classes": 1,
+  "torch_dtype": "float32",
+  "transformers_version": "4.29.2"
+}

configuration_gcn.py

@@ -0,0 +1,27 @@
+from transformers import PretrainedConfig
+
+class GCNConfig(PretrainedConfig):
+    model_type = "gcn"
+
+    def __init__(
+        self,
+        input_feature: int=64,
+        emb_input: int=20,
+        hidden_size: int=64,
+        n_layers: int=6,
+        num_classes: int=1,
+        **kwargs,
+    ):
+
+        self.input_feature = input_feature        # the dimension of input feature
+        self.emb_input = emb_input                # the embedding dimension of input feature
+        self.hidden_size = hidden_size            # the hidden size of GCN
+        self.n_layers = n_layers                  # the number of GCN layers
+        self.num_classes = num_classes            # the number of output classes
+
+        super().__init__(**kwargs)
+
+
+if __name__ == "__main__":
+    gcn_config = GCNConfig(input_feature=64, emb_input=20, hidden_size=64, n_layers=6, num_classes=1)
+    gcn_config.save_pretrained("custom-gcn")

modeling_gcn.py

@@ -0,0 +1,103 @@
+from torch_geometric.nn import GCNConv
+import torch.nn as nn
+import torch.nn.functional as F
+from torch_scatter import scatter
+from transformers import PreTrainedModel
+from gcn_model.configuration_gcn import GCNConfig
+import torch
+
+
+"""
+    MLP Layer used after graph vector representation
+"""
+class MLPReadout(nn.Module):
+
+    def __init__(self, input_dim, output_dim, L=2):  # L=nb_hidden_layers
+        super().__init__()
+        list_FC_layers = [nn.Linear(input_dim // 2 ** l, input_dim // 2 ** (l + 1), bias=True) for l in range(L)]
+        list_FC_layers.append(nn.Linear(input_dim // 2 ** L, output_dim, bias=True))
+        self.FC_layers = nn.ModuleList(list_FC_layers)
+        self.L = L
+
+    def forward(self, x):
+        y = x
+        for l in range(self.L):
+            y = self.FC_layers[l](y)
+            y = F.relu(y)
+        y = self.FC_layers[self.L](y)
+        return y
+
+class GCNNet(torch.nn.Module):
+    def __init__(self, input_feature=64, emb_input=20, hidden_size=64, n_layers=6, num_classes=1):
+        super(GCNNet, self).__init__()
+
+        self.embedding = torch.nn.Embedding(emb_input, hidden_size, padding_idx=0)
+        self.input_feature = input_feature
+        self.n_layers = n_layers  # 2层GCN
+        self.num_classes = num_classes
+
+        self.conv1 = GCNConv(hidden_size, hidden_size)
+
+        self.conv2 = GCNConv(hidden_size, 32)
+        self.mlp = MLPReadout(32, num_classes)
+
+    def forward(self, data):
+        x, edge_index, batch = data.x.long(), data.edge_index, data.batch
+        x = self.embedding(x.reshape(-1))
+
+        for i in range(self.n_layers):
+            x = F.relu(self.conv1(x, edge_index))
+
+        x = F.relu(self.conv2(x, edge_index))
+        x = scatter(x, batch, dim=-2, reduce='mean')
+        x = self.mlp(x)
+
+        return x.sequeeze(-1)
+
+class GCNModel(PreTrainedModel):
+    config_class = GCNConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.model = GCNNet(
+            input_feature=config.input_feature,
+            emb_input=config.emb_input,
+            hidden_size=config.hidden_size,
+            n_layers=config.n_layers,
+            num_classes=config.num_classes,
+        )
+
+    def forward(self, tensor):
+        return self.model.forward_features(tensor)
+
+# class GCNModelForMolecularPrediction(PreTrainedModel):
+#     config_class = GCNConfig
+#
+#     def __init__(self, config):
+#         super().__init__(config)
+#
+#         self.model = GCNNet(
+#             input_feature=config.input_feature,
+#             emb_input=config.emb_input,
+#             hidden_size=config.hidden_size,
+#             n_layers=config.n_layers,
+#             num_classes=config.num_classes,
+#         )
+#
+#     def forward(self, tensor):
+#         return self.model.forward_features(tensor)
+
+
+if __name__ == "__main__":
+    gcn_config = GCNConfig.from_pretrained("custom-gcn")
+
+    gcnd = GCNModel(gcn_config)
+    gcnd.model.load_state_dict(torch.load(r'G:\Trans_MXM\gcn_model\gcn.pt'))
+    gcnd.save_pretrained("custom-gcn")
+
+    # gcnd1 = GCNModelForMolecularPrediction(gcn_config)
+    #
+    # gcnd1.model.load_state_dict(torch.load(r'G:\Trans_MXM\gcn_model\gcn.pt'))
+    # gcnd1.save_pretrained("custom-gcn")
+

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b49c62eec2f337c8c36abe60088848da35409fd74ea35e0027f316ec92c1cc4f
+size 35716