[ { "path": "DCRN.py", "content": "import opt\nimport torch\nfrom torch import nn\nfrom torch.nn import Linear\nimport torch.nn.functional as F\nfrom torch.nn import Module, Parameter\n\n\n# AE encoder from DFCN\nclass AE_encoder(nn.Module):\n def __init__(self, ae_n_enc_1, ae_n_enc_2, ae_n_enc_3, n_input, n_z):\n super(AE_encoder, self).__init__()\n self.enc_1 = Linear(n_input, ae_n_enc_1)\n self.enc_2 = Linear(ae_n_enc_1, ae_n_enc_2)\n self.enc_3 = Linear(ae_n_enc_2, ae_n_enc_3)\n self.z_layer = Linear(ae_n_enc_3, n_z)\n self.act = nn.LeakyReLU(0.2, inplace=True)\n\n def forward(self, x):\n z = self.act(self.enc_1(x))\n z = self.act(self.enc_2(z))\n z = self.act(self.enc_3(z))\n z_ae = self.z_layer(z)\n return z_ae\n\n\n# AE decoder from DFCN\nclass AE_decoder(nn.Module):\n def __init__(self, ae_n_dec_1, ae_n_dec_2, ae_n_dec_3, n_input, n_z):\n super(AE_decoder, self).__init__()\n\n self.dec_1 = Linear(n_z, ae_n_dec_1)\n self.dec_2 = Linear(ae_n_dec_1, ae_n_dec_2)\n self.dec_3 = Linear(ae_n_dec_2, ae_n_dec_3)\n self.x_bar_layer = Linear(ae_n_dec_3, n_input)\n self.act = nn.LeakyReLU(0.2, inplace=True)\n\n def forward(self, z_ae):\n z = self.act(self.dec_1(z_ae))\n z = self.act(self.dec_2(z))\n z = self.act(self.dec_3(z))\n x_hat = self.x_bar_layer(z)\n return x_hat\n\n\n# Auto Encoder from DFCN\nclass AE(nn.Module):\n def __init__(self, ae_n_enc_1, ae_n_enc_2, ae_n_enc_3, ae_n_dec_1, ae_n_dec_2, ae_n_dec_3, n_input, n_z):\n super(AE, self).__init__()\n\n self.encoder = AE_encoder(\n ae_n_enc_1=ae_n_enc_1,\n ae_n_enc_2=ae_n_enc_2,\n ae_n_enc_3=ae_n_enc_3,\n n_input=n_input,\n n_z=n_z)\n\n self.decoder = AE_decoder(\n ae_n_dec_1=ae_n_dec_1,\n ae_n_dec_2=ae_n_dec_2,\n ae_n_dec_3=ae_n_dec_3,\n n_input=n_input,\n n_z=n_z)\n\n\n# GNNLayer from DFCN\nclass GNNLayer(Module):\n def __init__(self, in_features, out_features):\n super(GNNLayer, self).__init__()\n self.in_features = in_features\n self.out_features = out_features\n if opt.args.name == \"dblp\":\n self.act = nn.Tanh()\n self.weight = Parameter(torch.FloatTensor(out_features, in_features))\n else:\n self.act = nn.Tanh()\n self.weight = Parameter(torch.FloatTensor(in_features, out_features))\n torch.nn.init.xavier_uniform_(self.weight)\n\n def forward(self, features, adj, active=False):\n if active:\n if opt.args.name == \"dblp\":\n support = self.act(F.linear(features, self.weight))\n else:\n support = self.act(torch.mm(features, self.weight))\n else:\n if opt.args.name == \"dblp\":\n support = F.linear(features, self.weight)\n else:\n support = torch.mm(features, self.weight)\n output = torch.spmm(adj, support)\n az = torch.spmm(adj, output)\n return output, az\n\n\n# IGAE encoder from DFCN\nclass IGAE_encoder(nn.Module):\n def __init__(self, gae_n_enc_1, gae_n_enc_2, gae_n_enc_3, n_input):\n super(IGAE_encoder, self).__init__()\n self.gnn_1 = GNNLayer(n_input, gae_n_enc_1)\n self.gnn_2 = GNNLayer(gae_n_enc_1, gae_n_enc_2)\n self.gnn_3 = GNNLayer(gae_n_enc_2, gae_n_enc_3)\n self.s = nn.Sigmoid()\n\n def forward(self, x, adj):\n z_1, az_1 = self.gnn_1(x, adj, active=True)\n z_2, az_2 = self.gnn_2(z_1, adj, active=True)\n z_igae, az_3 = self.gnn_3(z_2, adj, active=False)\n z_igae_adj = self.s(torch.mm(z_igae, z_igae.t()))\n return z_igae, z_igae_adj, [az_1, az_2, az_3], [z_1, z_2, z_igae]\n\n\n# IGAE decoder from DFCN\nclass IGAE_decoder(nn.Module):\n def __init__(self, gae_n_dec_1, gae_n_dec_2, gae_n_dec_3, n_input):\n super(IGAE_decoder, self).__init__()\n self.gnn_4 = GNNLayer(gae_n_dec_1, gae_n_dec_2)\n self.gnn_5 = GNNLayer(gae_n_dec_2, gae_n_dec_3)\n self.gnn_6 = GNNLayer(gae_n_dec_3, n_input)\n self.s = nn.Sigmoid()\n\n def forward(self, z_igae, adj):\n z_1, az_1 = self.gnn_4(z_igae, adj, active=True)\n z_2, az_2 = self.gnn_5(z_1, adj, active=True)\n z_hat, az_3 = self.gnn_6(z_2, adj, active=True)\n z_hat_adj = self.s(torch.mm(z_hat, z_hat.t()))\n return z_hat, z_hat_adj, [az_1, az_2, az_3], [z_1, z_2, z_hat]\n\n\n# Improved Graph Auto Encoder from DFCN\nclass IGAE(nn.Module):\n def __init__(self, gae_n_enc_1, gae_n_enc_2, gae_n_enc_3, gae_n_dec_1, gae_n_dec_2, gae_n_dec_3, n_input):\n super(IGAE, self).__init__()\n # IGAE encoder\n self.encoder = IGAE_encoder(\n gae_n_enc_1=gae_n_enc_1,\n gae_n_enc_2=gae_n_enc_2,\n gae_n_enc_3=gae_n_enc_3,\n n_input=n_input)\n\n # IGAE decoder\n self.decoder = IGAE_decoder(\n gae_n_dec_1=gae_n_dec_1,\n gae_n_dec_2=gae_n_dec_2,\n gae_n_dec_3=gae_n_dec_3,\n n_input=n_input)\n\n\n# readout function\nclass Readout(nn.Module):\n def __init__(self, K):\n super(Readout, self).__init__()\n self.K = K\n\n def forward(self, Z):\n # calculate cluster-level embedding\n Z_tilde = []\n\n # step1: split the nodes into K groups\n # step2: average the node embedding in each group\n n_node = Z.shape[0]\n step = n_node // self.K\n for i in range(0, n_node, step):\n if n_node - i < 2 * step:\n Z_tilde.append(torch.mean(Z[i:n_node], dim=0))\n break\n else:\n Z_tilde.append(torch.mean(Z[i:i + step], dim=0))\n\n # the cluster-level embedding\n Z_tilde = torch.cat(Z_tilde, dim=0)\n return Z_tilde.view(1, -1)\n\n\n# Dual Correlation Reduction Network\nclass DCRN(nn.Module):\n def __init__(self, n_node=None):\n super(DCRN, self).__init__()\n\n # Auto Encoder\n self.ae = AE(\n ae_n_enc_1=opt.args.ae_n_enc_1,\n ae_n_enc_2=opt.args.ae_n_enc_2,\n ae_n_enc_3=opt.args.ae_n_enc_3,\n ae_n_dec_1=opt.args.ae_n_dec_1,\n ae_n_dec_2=opt.args.ae_n_dec_2,\n ae_n_dec_3=opt.args.ae_n_dec_3,\n n_input=opt.args.n_input,\n n_z=opt.args.n_z)\n\n # Improved Graph Auto Encoder From DFCN\n self.gae = IGAE(\n gae_n_enc_1=opt.args.gae_n_enc_1,\n gae_n_enc_2=opt.args.gae_n_enc_2,\n gae_n_enc_3=opt.args.gae_n_enc_3,\n gae_n_dec_1=opt.args.gae_n_dec_1,\n gae_n_dec_2=opt.args.gae_n_dec_2,\n gae_n_dec_3=opt.args.gae_n_dec_3,\n n_input=opt.args.n_input)\n\n # fusion parameter from DFCN\n self.a = Parameter(nn.init.constant_(torch.zeros(n_node, opt.args.n_z), 0.5), requires_grad=True)\n self.b = Parameter(nn.init.constant_(torch.zeros(n_node, opt.args.n_z), 0.5), requires_grad=True)\n self.alpha = Parameter(torch.zeros(1))\n\n # cluster layer (clustering assignment matrix)\n self.cluster_centers = Parameter(torch.Tensor(opt.args.n_clusters, opt.args.n_z), requires_grad=True)\n\n # readout function\n self.R = Readout(K=opt.args.n_clusters)\n\n # calculate the soft assignment distribution Q\n def q_distribute(self, Z, Z_ae, Z_igae):\n \"\"\"\n calculate the soft assignment distribution based on the embedding and the cluster centers\n Args:\n Z: fusion node embedding\n Z_ae: node embedding encoded by AE\n Z_igae: node embedding encoded by IGAE\n Returns:\n the soft assignment distribution Q\n \"\"\"\n q = 1.0 / (1.0 + torch.sum(torch.pow(Z.unsqueeze(1) - self.cluster_centers, 2), 2))\n q = (q.t() / torch.sum(q, 1)).t()\n\n q_ae = 1.0 / (1.0 + torch.sum(torch.pow(Z_ae.unsqueeze(1) - self.cluster_centers, 2), 2))\n q_ae = (q_ae.t() / torch.sum(q_ae, 1)).t()\n\n q_igae = 1.0 / (1.0 + torch.sum(torch.pow(Z_igae.unsqueeze(1) - self.cluster_centers, 2), 2))\n q_igae = (q_igae.t() / torch.sum(q_igae, 1)).t()\n\n return [q, q_ae, q_igae]\n\n def forward(self, X_tilde1, Am, X_tilde2, Ad):\n # node embedding encoded by AE\n Z_ae1 = self.ae.encoder(X_tilde1)\n Z_ae2 = self.ae.encoder(X_tilde2)\n\n # node embedding encoded by IGAE\n Z_igae1, A_igae1, AZ_1, Z_1 = self.gae.encoder(X_tilde1, Am)\n Z_igae2, A_igae2, AZ_2, Z_2 = self.gae.encoder(X_tilde2, Ad)\n\n # cluster-level embedding calculated by readout function\n Z_tilde_ae1 = self.R(Z_ae1)\n Z_tilde_ae2 = self.R(Z_ae2)\n Z_tilde_igae1 = self.R(Z_igae1)\n Z_tilde_igae2 = self.R(Z_igae2)\n\n # linear combination of view 1 and view 2\n Z_ae = (Z_ae1 + Z_ae2) / 2\n Z_igae = (Z_igae1 + Z_igae2) / 2\n\n # node embedding fusion from DFCN\n Z_i = self.a * Z_ae + self.b * Z_igae\n Z_l = torch.spmm(Am, Z_i)\n S = torch.mm(Z_l, Z_l.t())\n S = F.softmax(S, dim=1)\n Z_g = torch.mm(S, Z_l)\n Z = self.alpha * Z_g + Z_l\n\n # AE decoding\n X_hat = self.ae.decoder(Z)\n\n # IGAE decoding\n Z_hat, Z_adj_hat, AZ_de, Z_de = self.gae.decoder(Z, Am)\n sim = (A_igae1 + A_igae2) / 2\n A_hat = sim + Z_adj_hat\n\n # node embedding and cluster-level embedding\n Z_ae_all = [Z_ae1, Z_ae2, Z_tilde_ae1, Z_tilde_ae2]\n Z_gae_all = [Z_igae1, Z_igae2, Z_tilde_igae1, Z_tilde_igae2]\n\n # the soft assignment distribution Q\n Q = self.q_distribute(Z, Z_ae, Z_igae)\n\n # propagated embedding AZ_all and embedding Z_all\n AZ_en = []\n Z_en = []\n for i in range(len(AZ_1)):\n AZ_en.append((AZ_1[i]+AZ_2[i])/2)\n Z_en.append((Z_1[i]+Z_2[i])/2)\n AZ_all = [AZ_en, AZ_de]\n Z_all = [Z_en, Z_de]\n\n return X_hat, Z_hat, A_hat, sim, Z_ae_all, Z_gae_all, Q, Z, AZ_all, Z_all\n" }, { "path": "LICENSE", "content": "MIT License\n\nCopyright (c) 2021 yueliu1999\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n" }, { "path": "README.md", "content": "[stars-img]: https://img.shields.io/github/stars/yueliu1999/DCRN?color=yellow\n[stars-url]: https://github.com/yueliu1999/DCRN/stargazers\n[fork-img]: https://img.shields.io/github/forks/yueliu1999/DCRN?color=lightblue&label=fork\n[fork-url]: https://github.com/yueliu1999/DCRN/network/members\n[visitors-img]: https://visitor-badge.glitch.me/badge?page_id=yueliu1999.DCRN\n[adgc-url]: https://github.com/yueliu1999/DCRN\n\n# Dual Correlation Reduction Network\n\n

\n \n \n \n \n

\n\n\n\n[![GitHub stars][stars-img]][stars-url]\n[![GitHub forks][fork-img]][fork-url]\n[![visitors][visitors-img]][adgc-url]\n\n\nAn official source code for paper [Deep Graph Clustering via Dual Correlation Reduction](https://www.researchgate.net/profile/Yue-Liu-240/publication/357271184_Deep_Graph_Clustering_via_Dual_Correlation_Reduction/links/61c466e68bb20101842f9a92/Deep-Graph-Clustering-via-Dual-Correlation-Reduction.pdf), accepted by AAAI 2022. Any communications or issues are welcomed. Please contact yueliu19990731@163.com. If you find this repository useful to your research or work, it is really appreciate to star this repository. :heart:\n\n-------------\n\n### Overview\n\n

\n Deep graph clustering, which aims to reveal the underlying graph structure and divide the nodes into different groups, has attracted intensive attention in recent years. However, we observe that, in the process of node encoding, existing methods suffer from representation collapse which tends to map all data into a same representation. Consequently, the discriminative capability of node representations is limited, leading to unsatisfied clustering performance. To address this issue, we propose a novel self-supervised deep graph clustering method termed Dual Correlation Reduction Network (DCRN) by reducing information correlation in a dual manner. Specifically, in our method, we first design a siamese network to encode samples. Then by forcing the cross-view sample correlation matrix and cross-view feature correlation matrix to approximate two identity matrices, respectively, we reduce the information correlation in dual level, thus improve the discriminative capability of the resulting features. Moreover, in order to alleviate representation collapse caused by over-smoothing in GCN, we introduce a propagation-regularization term to enable the network to gain long-distance information with shallow network structure. Extensive experimental results on six benchmark datasets demonstrate the effectiveness of the proposed DCRN against the existing state-of-the-art methods.\n

\n
\n \n
\n\n\n
\n Illustration of the Dual Correlation Reduction Network (DCRN).\n
\n\n\n### Requirements\n\nThe proposed DCRN is implemented with python 3.8.5 on a NVIDIA 3090 GPU. \n\nPython package information is summarized in **requirements.txt**:\n\n- torch==1.8.0\n- tqdm==4.50.2\n- numpy==1.19.2\n- munkres==1.1.4\n- scikit_learn==1.0.1\n\n### Pre-training\nWe release the pre-training code.\n\n- Google Drive: [Link](https://drive.google.com/file/d/1XRlu3Ahgwin52jluqFu2aBW6wjCwjY4M/view?usp=sharing)\n- Nut store: [Link](https://www.jianguoyun.com/p/DXCOQEYQwdaSChiEjrsEIAA)\n\n### Quick Start\n\n- Step1: use the **dblp.zip** file or download other datasets from [Awesome Deep Graph Clustering/Benchmark Datasets](https://github.com/yueliu1999/Awesome-Deep-Graph-Clustering#benchmark-datasets) \n\n- Step2: unzip the dataset into the **./dataset** folder\n\n- Step3: run \n\n ```\n python main.py --name dblp --seed 3 --alpha_value 0.2 --lambda_value 10 --gamma_value 1e3 --lr 1e-4\n ```\n\nParameter setting\n\n- name: the name of dataset\n- seed: the random seed. 10 runs under different random seeds.\n- alpha_value: the teleport probability in graph diffusion\n - PUBMED: 0.1\n - DBLP, CITE, ACM, AMAP, CORAFULL: 0.2\n- lambda_value: the coefficient of clustering guidance loss.\n - all datasets: 10\n- gamma_value: the coefficient of propagation regularization\n - all datasets: 1e3\n- lr: learning rate\n - DBLP 1e-4\n - ACM: 5e-5\n - AMAP: 1e-3\n - CITE, PUBMED, CORAFULL: 1e-5\n\n\n\nTips: Limited by the GPU memory, PUBMED and CORAFULL might be out of memory during training. Thus, we adpot batch training on PUBMED and CORAFULL dataseets and the batch size is set to 2000. Please use the batch training version of DCRN [here](https://drive.google.com/file/d/185GLObsQQL3Y-dQ2aIin5YrXuA-dgpnU/view?usp=sharing).\n\n\n\n### Results\n\n
\n \n
\n\n\n\n
\n \n
\n\n\n### Citation\n\nIf you use code or datasets in this repository for your research, please cite our paper.\n\n```\n@inproceedings{DCRN,\n title={Deep Graph Clustering via Dual Correlation Reduction},\n author={Liu, Yue and Tu, Wenxuan and Zhou, Sihang and Liu, Xinwang and Song, Linxuan and Yang, Xihong and Zhu, En},\n booktitle={Proceedings of the AAAI Conference on Artificial Intelligence},\n volume={36},\n number={7},\n pages={7603-7611},\n year={2022}\n}\n\n@article{Deep_graph_clustering_survey,\n author = {Liu, Yue and Xia, Jun and Zhou, Sihang and Wang, Siwei and Guo, Xifeng and Yang, Xihong and Liang, Ke and Tu, Wenxuan and Li, Z. Stan and Liu, Xinwang},\n journal = {arXiv preprint arXiv:2211.12875},\n title = {A Survey of Deep Graph Clustering: Taxonomy, Challenge, and Application},\n year = {2022}\n}\n```\n\n\n" }, { "path": "main.py", "content": "from train import *\nfrom DCRN import DCRN\n\n\nif __name__ == '__main__':\n # setup\n setup()\n\n # data pre-precessing: X, y, A, A_norm, Ad\n X, y, A = load_graph_data(opt.args.name, show_details=False)\n A_norm = normalize_adj(A, self_loop=True, symmetry=False)\n Ad = diffusion_adj(A, mode=\"ppr\", transport_rate=opt.args.alpha_value)\n\n # to torch tensor\n X = numpy_to_torch(X).to(opt.args.device)\n A_norm = numpy_to_torch(A_norm, sparse=True).to(opt.args.device)\n Ad = numpy_to_torch(Ad).to(opt.args.device)\n\n # Dual Correlation Reduction Network\n model = DCRN(n_node=X.shape[0]).to(opt.args.device)\n\n # deep graph clustering\n acc, nmi, ari, f1 = train(model, X, y, A, A_norm, Ad)\n print(\"ACC: {:.4f},\".format(acc), \"NMI: {:.4f},\".format(nmi), \"ARI: {:.4f},\".format(ari), \"F1: {:.4f}\".format(f1))\n" }, { "path": "opt.py", "content": "import argparse\n\nparser = argparse.ArgumentParser(description='DCRN', formatter_class=argparse.ArgumentDefaultsHelpFormatter)\n\n# setting\nparser.add_argument('--name', type=str, default=\"dblp\")\nparser.add_argument('--cuda', type=bool, default=True)\nparser.add_argument('--seed', type=int, default=3)\nparser.add_argument('--alpha_value', type=float, default=0.2)\nparser.add_argument('--lambda_value', type=float, default=10)\nparser.add_argument('--gamma_value', type=float, default=1e3)\nparser.add_argument('--lr', type=float, default=1e-4)\nparser.add_argument('--n_z', type=int, default=20)\nparser.add_argument('--epoch', type=int, default=400)\nparser.add_argument('--show_training_details', type=bool, default=False)\n\n\n# AE structure parameter from DFCN\nparser.add_argument('--ae_n_enc_1', type=int, default=128)\nparser.add_argument('--ae_n_enc_2', type=int, default=256)\nparser.add_argument('--ae_n_enc_3', type=int, default=512)\nparser.add_argument('--ae_n_dec_1', type=int, default=512)\nparser.add_argument('--ae_n_dec_2', type=int, default=256)\nparser.add_argument('--ae_n_dec_3', type=int, default=128)\n\n# IGAE structure parameter from DFCN\nparser.add_argument('--gae_n_enc_1', type=int, default=128)\nparser.add_argument('--gae_n_enc_2', type=int, default=256)\nparser.add_argument('--gae_n_enc_3', type=int, default=20)\nparser.add_argument('--gae_n_dec_1', type=int, default=20)\nparser.add_argument('--gae_n_dec_2', type=int, default=256)\nparser.add_argument('--gae_n_dec_3', type=int, default=128)\n\n# clustering performance: acc, nmi, ari, f1\nparser.add_argument('--acc', type=float, default=0)\nparser.add_argument('--nmi', type=float, default=0)\nparser.add_argument('--ari', type=float, default=0)\nparser.add_argument('--f1', type=float, default=0)\n\nargs = parser.parse_args()\n" }, { "path": "requirements.txt", "content": "torch==1.8.0\ntqdm==4.50.2\nnumpy==1.22.0\nmunkres==1.1.4\nscikit_learn==1.0.1\n\n" }, { "path": "train.py", "content": "import tqdm\nfrom utils import *\nfrom torch.optim import Adam\n\n\ndef train(model, X, y, A, A_norm, Ad):\n \"\"\"\n train our model\n Args:\n model: Dual Correlation Reduction Network\n X: input feature matrix\n y: input label\n A: input origin adj\n A_norm: normalized adj\n Ad: graph diffusion\n Returns: acc, nmi, ari, f1\n \"\"\"\n print(\"Training…\")\n # calculate embedding similarity and cluster centers\n sim, centers = model_init(model, X, y, A_norm)\n\n # initialize cluster centers\n model.cluster_centers.data = torch.tensor(centers).to(opt.args.device)\n\n # edge-masked adjacency matrix (Am): remove edges based on feature-similarity\n Am = remove_edge(A, sim, remove_rate=0.1)\n\n optimizer = Adam(model.parameters(), lr=opt.args.lr)\n for epoch in tqdm.tqdm(range(opt.args.epoch)):\n # add gaussian noise to X\n X_tilde1, X_tilde2 = gaussian_noised_feature(X)\n\n # input & output\n X_hat, Z_hat, A_hat, _, Z_ae_all, Z_gae_all, Q, Z, AZ_all, Z_all = model(X_tilde1, Ad, X_tilde2, Am)\n\n # calculate loss: L_{DICR}, L_{REC} and L_{KL}\n L_DICR = dicr_loss(Z_ae_all, Z_gae_all, AZ_all, Z_all)\n L_REC = reconstruction_loss(X, A_norm, X_hat, Z_hat, A_hat)\n L_KL = distribution_loss(Q, target_distribution(Q[0].data))\n loss = L_DICR + L_REC + opt.args.lambda_value * L_KL\n\n # optimization\n optimizer.zero_grad()\n loss.backward(retain_graph=True)\n optimizer.step()\n\n # clustering & evaluation\n acc, nmi, ari, f1, _ = clustering(Z, y)\n if acc > opt.args.acc:\n opt.args.acc = acc\n opt.args.nmi = nmi\n opt.args.ari = ari\n opt.args.f1 = f1\n\n return opt.args.acc, opt.args.nmi, opt.args.ari, opt.args.f1\n" }, { "path": "utils.py", "content": "import opt\nimport torch\nimport random\nimport numpy as np\nfrom sklearn import metrics\nfrom munkres import Munkres\nimport torch.nn.functional as F\nfrom sklearn.cluster import KMeans\nfrom sklearn.decomposition import PCA\nfrom sklearn.metrics import adjusted_rand_score as ari_score\nfrom sklearn.metrics.cluster import normalized_mutual_info_score as nmi_score\n\n\ndef setup():\n \"\"\"\n setup\n - name: the name of dataset\n - device: CPU / GPU\n - seed: random seed\n - n_clusters: num of cluster\n - n_input: dimension of feature\n - alpha_value: alpha value for graph diffusion\n - lambda_value: lambda value for clustering guidance\n - gamma_value: gamma value for propagation regularization\n - lr: learning rate\n Return: None\n\n \"\"\"\n print(\"setting:\")\n setup_seed(opt.args.seed)\n if opt.args.name == 'acm':\n opt.args.n_clusters = 3\n opt.args.n_input = 100\n opt.args.alpha_value = 0.2\n opt.args.lambda_value = 10\n opt.args.gamma_value = 1e3\n opt.args.lr = 5e-5\n\n elif opt.args.name == 'dblp':\n opt.args.n_clusters = 4\n opt.args.n_input = 50\n opt.args.alpha_value = 0.2\n opt.args.lambda_value = 10\n opt.args.gamma_value = 1e3\n opt.args.lr = 1e-4\n\n elif opt.args.name == 'cite':\n opt.args.n_clusters = 6\n opt.args.n_input = 100\n opt.args.alpha_value = 0.2\n opt.args.lambda_value = 10\n opt.args.gamma_value = 1e3\n opt.args.lr = 1e-5\n\n elif opt.args.name == 'amap':\n opt.args.n_clusters = 8\n opt.args.n_input = 100\n opt.args.alpha_value = 0.2\n opt.args.lambda_value = 10\n opt.args.gamma_value = 1e3\n opt.args.lr = 1e-3\n\n else:\n print(\"error!\")\n print(\"please add the new dataset's parameters\")\n print(\"------------------------------\")\n print(\"dataset : \")\n print(\"device : \")\n print(\"random seed : \")\n print(\"clusters : \")\n print(\"alpha value : \")\n print(\"lambda value : \")\n print(\"gamma value : \")\n print(\"learning rate : \")\n print(\"------------------------------\")\n exit(0)\n\n opt.args.device = torch.device(\"cuda\" if opt.args.cuda else \"cpu\")\n print(\"------------------------------\")\n print(\"dataset : {}\".format(opt.args.name))\n print(\"device : {}\".format(opt.args.device))\n print(\"random seed : {}\".format(opt.args.seed))\n print(\"clusters : {}\".format(opt.args.n_clusters))\n print(\"alpha value : {}\".format(opt.args.alpha_value))\n print(\"lambda value : {}\".format(opt.args.lambda_value))\n print(\"gamma value : {:.0e}\".format(opt.args.gamma_value))\n print(\"learning rate : {:.0e}\".format(opt.args.lr))\n print(\"------------------------------\")\n\n\ndef setup_seed(seed):\n \"\"\"\n setup random seed to fix the result\n Args:\n seed: random seed\n Returns: None\n \"\"\"\n torch.manual_seed(seed)\n torch.cuda.manual_seed(seed)\n torch.cuda.manual_seed_all(seed)\n np.random.seed(seed)\n random.seed(seed)\n torch.manual_seed(seed)\n torch.backends.cudnn.benchmark = False\n torch.backends.cudnn.deterministic = True\n\n\ndef numpy_to_torch(a, sparse=False):\n \"\"\"\n numpy array to torch tensor\n :param a: the numpy array\n :param sparse: is sparse tensor or not\n :return: torch tensor\n \"\"\"\n if sparse:\n a = torch.sparse.Tensor(a)\n a = a.to_sparse()\n else:\n a = torch.FloatTensor(a)\n return a\n\n\ndef torch_to_numpy(t):\n \"\"\"\n torch tensor to numpy array\n :param t: the torch tensor\n :return: numpy array\n \"\"\"\n return t.numpy()\n\n\ndef load_graph_data(dataset_name, show_details=False):\n \"\"\"\n load graph data\n :param dataset_name: the name of the dataset\n :param show_details: if show the details of dataset\n - dataset name\n - features' shape\n - labels' shape\n - adj shape\n - edge num\n - category num\n - category distribution\n :return: the features, labels and adj\n \"\"\"\n load_path = \"dataset/\" + dataset_name + \"/\" + dataset_name\n feat = np.load(load_path+\"_feat.npy\", allow_pickle=True)\n label = np.load(load_path+\"_label.npy\", allow_pickle=True)\n adj = np.load(load_path+\"_adj.npy\", allow_pickle=True)\n if show_details:\n print(\"++++++++++++++++++++++++++++++\")\n print(\"---details of graph dataset---\")\n print(\"++++++++++++++++++++++++++++++\")\n print(\"dataset name: \", dataset_name)\n print(\"feature shape: \", feat.shape)\n print(\"label shape: \", label.shape)\n print(\"adj shape: \", adj.shape)\n print(\"undirected edge num: \", int(np.nonzero(adj)[0].shape[0]/2))\n print(\"category num: \", max(label)-min(label)+1)\n print(\"category distribution: \")\n for i in range(max(label)+1):\n print(\"label\", i, end=\":\")\n print(len(label[np.where(label == i)]))\n print(\"++++++++++++++++++++++++++++++\")\n\n # X pre-processing\n pca = PCA(n_components=opt.args.n_input)\n feat = pca.fit_transform(feat)\n return feat, label, adj\n\n\ndef normalize_adj(adj, self_loop=True, symmetry=False):\n \"\"\"\n normalize the adj matrix\n :param adj: input adj matrix\n :param self_loop: if add the self loop or not\n :param symmetry: symmetry normalize or not\n :return: the normalized adj matrix\n \"\"\"\n # add the self_loop\n if self_loop:\n adj_tmp = adj + np.eye(adj.shape[0])\n else:\n adj_tmp = adj\n\n # calculate degree matrix and it's inverse matrix\n d = np.diag(adj_tmp.sum(0))\n d_inv = np.linalg.inv(d)\n\n # symmetry normalize: D^{-0.5} A D^{-0.5}\n if symmetry:\n sqrt_d_inv = np.sqrt(d_inv)\n norm_adj = np.matmul(np.matmul(sqrt_d_inv, adj_tmp), adj_tmp)\n\n # non-symmetry normalize: D^{-1} A\n else:\n norm_adj = np.matmul(d_inv, adj_tmp)\n\n return norm_adj\n\n\ndef gaussian_noised_feature(X):\n \"\"\"\n add gaussian noise to the attribute matrix X\n Args:\n X: the attribute matrix\n Returns: the noised attribute matrix X_tilde\n \"\"\"\n N_1 = torch.Tensor(np.random.normal(1, 0.1, X.shape)).to(opt.args.device)\n N_2 = torch.Tensor(np.random.normal(1, 0.1, X.shape)).to(opt.args.device)\n X_tilde1 = X * N_1\n X_tilde2 = X * N_2\n return X_tilde1, X_tilde2\n\n\ndef diffusion_adj(adj, mode=\"ppr\", transport_rate=0.2):\n \"\"\"\n graph diffusion\n :param adj: input adj matrix\n :param mode: the mode of graph diffusion\n :param transport_rate: the transport rate\n - personalized page rank\n -\n :return: the graph diffusion\n \"\"\"\n # add the self_loop\n adj_tmp = adj + np.eye(adj.shape[0])\n\n # calculate degree matrix and it's inverse matrix\n d = np.diag(adj_tmp.sum(0))\n d_inv = np.linalg.inv(d)\n sqrt_d_inv = np.sqrt(d_inv)\n\n # calculate norm adj\n norm_adj = np.matmul(np.matmul(sqrt_d_inv, adj_tmp), sqrt_d_inv)\n\n # calculate graph diffusion\n if mode == \"ppr\":\n diff_adj = transport_rate * np.linalg.inv((np.eye(d.shape[0]) - (1 - transport_rate) * norm_adj))\n\n return diff_adj\n\n\ndef remove_edge(A, similarity, remove_rate=0.1):\n \"\"\"\n remove edge based on embedding similarity\n Args:\n A: the origin adjacency matrix\n similarity: cosine similarity matrix of embedding\n remove_rate: the rate of removing linkage relation\n Returns:\n Am: edge-masked adjacency matrix\n \"\"\"\n # remove edges based on cosine similarity of embedding\n n_node = A.shape[0]\n for i in range(n_node):\n A[i, torch.argsort(similarity[i].cpu())[:int(round(remove_rate * n_node))]] = 0\n\n # normalize adj\n Am = normalize_adj(A, self_loop=True, symmetry=False)\n Am = numpy_to_torch(Am).to(opt.args.device)\n return Am\n\n\ndef load_pretrain_parameter(model):\n \"\"\"\n load pretrained parameters\n Args:\n model: Dual Correlation Reduction Network\n Returns: model\n \"\"\"\n pretrained_dict = torch.load('model_pretrain/{}_pretrain.pkl'.format(opt.args.name), map_location='cpu')\n model_dict = model.state_dict()\n pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}\n model_dict.update(pretrained_dict)\n model.load_state_dict(model_dict)\n return model\n\n\ndef model_init(model, X, y, A_norm):\n \"\"\"\n load the pre-train model and calculate similarity and cluster centers\n Args:\n model: Dual Correlation Reduction Network\n X: input feature matrix\n y: input label\n A_norm: normalized adj\n Returns: embedding similarity matrix\n \"\"\"\n # load pre-train model\n model = load_pretrain_parameter(model)\n\n # calculate embedding similarity\n with torch.no_grad():\n _, _, _, sim, _, _, _, Z, _, _ = model(X, A_norm, X, A_norm)\n\n # calculate cluster centers\n acc, nmi, ari, f1, centers = clustering(Z, y)\n\n return sim, centers\n\n\n# the reconstruction function from DFCN\ndef reconstruction_loss(X, A_norm, X_hat, Z_hat, A_hat):\n \"\"\"\n reconstruction loss L_{}\n Args:\n X: the origin feature matrix\n A_norm: the normalized adj\n X_hat: the reconstructed X\n Z_hat: the reconstructed Z\n A_hat: the reconstructed A\n Returns: the reconstruction loss\n \"\"\"\n loss_ae = F.mse_loss(X_hat, X)\n loss_w = F.mse_loss(Z_hat, torch.spmm(A_norm, X))\n loss_a = F.mse_loss(A_hat, A_norm.to_dense())\n loss_igae = loss_w + 0.1 * loss_a\n loss_rec = loss_ae + loss_igae\n return loss_rec\n\n\ndef target_distribution(Q):\n \"\"\"\n calculate the target distribution (student-t distribution)\n Args:\n Q: the soft assignment distribution\n Returns: target distribution P\n \"\"\"\n weight = Q ** 2 / Q.sum(0)\n P = (weight.t() / weight.sum(1)).t()\n return P\n\n\n# clustering guidance from DFCN\ndef distribution_loss(Q, P):\n \"\"\"\n calculate the clustering guidance loss L_{KL}\n Args:\n Q: the soft assignment distribution\n P: the target distribution\n Returns: L_{KL}\n \"\"\"\n loss = F.kl_div((Q[0].log() + Q[1].log() + Q[2].log()) / 3, P, reduction='batchmean')\n return loss\n\n\ndef r_loss(AZ, Z):\n \"\"\"\n the loss of propagated regularization (L_R)\n Args:\n AZ: the propagated embedding\n Z: embedding\n Returns: L_R\n \"\"\"\n loss = 0\n for i in range(2):\n for j in range(3):\n p_output = F.softmax(AZ[i][j], dim=1)\n q_output = F.softmax(Z[i][j], dim=1)\n log_mean_output = ((p_output + q_output) / 2).log()\n loss += (F.kl_div(log_mean_output, p_output, reduction='batchmean') +\n F.kl_div(log_mean_output, p_output, reduction='batchmean')) / 2\n return loss\n\n\ndef off_diagonal(x):\n \"\"\"\n off-diagonal elements of x\n Args:\n x: the input matrix\n Returns: the off-diagonal elements of x\n \"\"\"\n n, m = x.shape\n assert n == m\n return x.flatten()[:-1].view(n - 1, n + 1)[:, 1:].flatten()\n\n\ndef cross_correlation(Z_v1, Z_v2):\n \"\"\"\n calculate the cross-view correlation matrix S\n Args:\n Z_v1: the first view embedding\n Z_v2: the second view embedding\n Returns: S\n \"\"\"\n return torch.mm(F.normalize(Z_v1, dim=1), F.normalize(Z_v2, dim=1).t())\n\n\ndef correlation_reduction_loss(S):\n \"\"\"\n the correlation reduction loss L: MSE for S and I (identical matrix)\n Args:\n S: the cross-view correlation matrix S\n Returns: L\n \"\"\"\n return torch.diagonal(S).add(-1).pow(2).mean() + off_diagonal(S).pow(2).mean()\n\n\ndef dicr_loss(Z_ae, Z_igae, AZ, Z):\n \"\"\"\n Dual Information Correlation Reduction loss L_{DICR}\n Args:\n Z_ae: AE embedding including two-view node embedding [0, 1] and two-view cluster-level embedding [2, 3]\n Z_igae: IGAE embedding including two-view node embedding [0, 1] and two-view cluster-level embedding [2, 3]\n AZ: the propagated fusion embedding AZ\n Z: the fusion embedding Z\n Returns:\n L_{DICR}\n \"\"\"\n # Sample-level Correlation Reduction (SCR)\n # cross-view sample correlation matrix\n S_N_ae = cross_correlation(Z_ae[0], Z_ae[1])\n S_N_igae = cross_correlation(Z_igae[0], Z_igae[1])\n # loss of SCR\n L_N_ae = correlation_reduction_loss(S_N_ae)\n L_N_igae = correlation_reduction_loss(S_N_igae)\n\n # Feature-level Correlation Reduction (FCR)\n # cross-view feature correlation matrix\n S_F_ae = cross_correlation(Z_ae[2].t(), Z_ae[3].t())\n S_F_igae = cross_correlation(Z_igae[2].t(), Z_igae[3].t())\n\n # loss of FCR\n L_F_ae = correlation_reduction_loss(S_F_ae)\n L_F_igae = correlation_reduction_loss(S_F_igae)\n\n if opt.args.name == \"dblp\" or opt.args.name == \"acm\":\n L_N = 0.01 * L_N_ae + 10 * L_N_igae\n L_F = 0.5 * L_F_ae + 0.5 * L_F_igae\n else:\n L_N = 0.1 * L_N_ae + 5 * L_N_igae\n L_F = L_F_ae + L_F_igae\n\n # propagated regularization\n L_R = r_loss(AZ, Z)\n\n # loss of DICR\n loss_dicr = L_N + L_F + opt.args.gamma_value * L_R\n\n return loss_dicr\n\n\ndef clustering(Z, y):\n \"\"\"\n clustering based on embedding\n Args:\n Z: the input embedding\n y: the ground truth\n\n Returns: acc, nmi, ari, f1, clustering centers\n \"\"\"\n model = KMeans(n_clusters=opt.args.n_clusters, n_init=20)\n cluster_id = model.fit_predict(Z.data.cpu().numpy())\n acc, nmi, ari, f1 = eva(y, cluster_id, show_details=opt.args.show_training_details)\n return acc, nmi, ari, f1, model.cluster_centers_\n\n\ndef cluster_acc(y_true, y_pred):\n \"\"\"\n calculate clustering acc and f1-score\n Args:\n y_true: the ground truth\n y_pred: the clustering id\n\n Returns: acc and f1-score\n \"\"\"\n y_true = y_true - np.min(y_true)\n l1 = list(set(y_true))\n num_class1 = len(l1)\n l2 = list(set(y_pred))\n num_class2 = len(l2)\n ind = 0\n if num_class1 != num_class2:\n for i in l1:\n if i in l2:\n pass\n else:\n y_pred[ind] = i\n ind += 1\n l2 = list(set(y_pred))\n numclass2 = len(l2)\n if num_class1 != numclass2:\n print('error')\n return\n cost = np.zeros((num_class1, numclass2), dtype=int)\n for i, c1 in enumerate(l1):\n mps = [i1 for i1, e1 in enumerate(y_true) if e1 == c1]\n for j, c2 in enumerate(l2):\n mps_d = [i1 for i1 in mps if y_pred[i1] == c2]\n cost[i][j] = len(mps_d)\n m = Munkres()\n cost = cost.__neg__().tolist()\n indexes = m.compute(cost)\n new_predict = np.zeros(len(y_pred))\n for i, c in enumerate(l1):\n c2 = l2[indexes[i][1]]\n ai = [ind for ind, elm in enumerate(y_pred) if elm == c2]\n new_predict[ai] = c\n acc = metrics.accuracy_score(y_true, new_predict)\n f1_macro = metrics.f1_score(y_true, new_predict, average='macro')\n return acc, f1_macro\n\n\ndef eva(y_true, y_pred, show_details=True):\n \"\"\"\n evaluate the clustering performance\n Args:\n y_true: the ground truth\n y_pred: the predicted label\n show_details: if print the details\n Returns: None\n \"\"\"\n acc, f1 = cluster_acc(y_true, y_pred)\n nmi = nmi_score(y_true, y_pred, average_method='arithmetic')\n ari = ari_score(y_true, y_pred)\n if show_details:\n print(':acc {:.4f}'.format(acc), ', nmi {:.4f}'.format(nmi), ', ari {:.4f}'.format(ari),\n ', f1 {:.4f}'.format(f1))\n return acc, nmi, ari, f1\n" } ]