├── LICENSE ├── README.md ├── data ├── aminer_data │ └── __init__.py ├── github_data │ └── __init__.py ├── instagram_data │ └── __ └── yelp_data │ └── __init__.py ├── finetune.py ├── finetune ├── aminer_data │ └── finetune_feature.txt ├── github_data │ └── __init__.py ├── instagram_data │ └── __init__.py └── yelp_data │ └── __init__.py ├── main_graph_image_gcl.py ├── main_graph_text_gcl.py ├── pretrain └── aminer.pt ├── requirements.txt ├── scripts ├── CMCL.py ├── modules_layer.py ├── modules_model.py └── moodules_graph.py └── utils ├── __init__.py ├── focalloss.py ├── getdata.py ├── params.py └── util.py /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 Meta-HG 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # CM-GCL 2 | Co-Modality Graph Contrastive Learning for Imbalanced Node Classification 3 | ==== 4 | Official source code of "Co-Modality Graph Contrastive Learning for Imbalanced Node Classification 5 | " (NeurIPS 2022 https://openreview.net/pdf?id=f_kvHrM4Q0). 6 | 7 | ## Requirements 8 | 9 | This code is developed and tested with python 3.8.10. and the required packages are listed in the `requirements.txt`. 10 | 11 | Please run `pip install -r requirements.txt` to install all the dependencies. 12 | 13 | ## Usage 14 | ### Data Download 15 | As the data size is too large, please click on the link to download the [AMiner Data](https://drive.google.com/file/d/1TJfviH4--HPp12jQnqbJEpO58GrHCexJ/view?usp=share_link). 16 | 17 | Tips: the file size is quite big including (a.) the AMiner data, (b.) the fine-tuned node features (finetune_feature_20301141240.txt), and (c.) the pre-trained pt model file (aminer_node_text_20301141240.pt). 18 | 19 | Once finish downloading, please unzip the file and further put the unzipped data (part a) to the "data/aminer_data/" folder; the fine-tuned node features (part b) to "finetune/aminer_data/" folder; 20 | the pre-trained pt file (part c) into the "pretrain/" folder; 21 | 22 | ### Model Pre-training 23 | For co-modality pre-training, we consider two types of modality combinations, i.e., graph modality and text modality, and 24 | graph modality and image modality in our paper: 25 | 26 | ```main_graph_text_gcl.py``` contains the code of graph contrastive learning over real-world datasets including raw text information and graph data. 27 | 28 | ```main_graph_image_gcl.py``` contains the code of graph contrastive learning over real-world datasets including raw images and graph data. 29 | 30 | ### Model Fine-tuning 31 | 32 | ```finetune.py``` contains the code of model finetuning for downstream tasks. 33 | 34 | The default setting for AMiner data is all set. If you want to train the model, please run the code ```python main_graph_text_gcl.py``` over AMiner data. 35 | It may take a while, you can also simply run the code ```python finetune.py``` to reproduce our results. We also provide a sample running log below. 36 | 37 | 38 | ## Dataset 39 | 40 | AMiner data is a paper-citation academic network containing the raw text content. The dataset of this paper is already provided to train our model. Please feel free to refer to the website for more details https://www.aminer.cn/aminer_data. 41 | 42 | Yelp data is a review social graph containg the raw text content. Please refer to the website for more details http://odds.cs.stonybrook.edu/yelpchi-dataset/. If you need the raw text data for your research purpose, please email to the author listed on the website. 43 | 44 | GitHub data is a graph for detecting malicious repository on social coding platform. Since our model uses raw text data for model training, it may cause privacy issues if we public the data. 45 | 46 | Instagram data is a social graph including raw text and image data for detecting drug traffickers on social platform. As our model needs raw image and text data for model training, it may cause privacy issues. 47 | If you want to implement our code (the combination of graph modality and image modality) on your datasets, please follow the format we provided in the instagram_data folder. 48 | 49 | 50 | 51 | ## Contact 52 | Yiyue Qian - yqian5@nd.edu or yxq250@case.edu 53 | 54 | Discussions, suggestions and questions are always welcome! 55 | 56 | 57 | 58 | ## Citation 59 | If our work helps you, please cite our paper: 60 | 61 | ``` 62 | @inproceedings{qianco, 63 | title={Co-Modality Graph Contrastive Learning for Imbalanced Node Classification}, 64 | author={Qian, Yiyue and Zhang, Chunhui and Zhang, Yiming and Wen, Qianlong and Ye, Yanfang and Zhang, Chuxu}, 65 | booktitle={Advances in Neural Information Processing Systems}, 66 | year={2022} 67 | } 68 | ``` 69 | 70 | ## Logger 71 | This is a sample running logger which records the output and the model performance for AMiner data (also provide a output.log in aminer_data folder): 72 | ``` 73 | python main_graph_text_gcl.py 74 | 75 | Epoch: 1 76 | 100%|██████████| 145/145 [00:23<00:00, 6.19it/s, lr=0.001, train_loss=41.7] 77 | 100%|██████████| 37/37 [00:07<00:00, 5.08it/s, valid_loss=10.6] 78 | 0%| | 0/145 [00:00 ' \ 46 | + str(self.out_features) + ')' 47 | 48 | 49 | class GraphAttentionLayer(nn.Module): 50 | """ 51 | Simple GAT layer, similar to https://arxiv.org/abs/1710.10903 52 | """ 53 | def __init__(self, in_features, out_features, dropout, alpha, concat=True): 54 | super(GraphAttentionLayer, self).__init__() 55 | self.dropout = dropout 56 | self.in_features = in_features 57 | self.out_features = out_features 58 | self.alpha = alpha 59 | self.concat = concat 60 | 61 | self.W = nn.Parameter(torch.empty(size=(in_features, out_features))) 62 | nn.init.xavier_uniform_(self.W.data, gain=1.414) 63 | self.a = nn.Parameter(torch.empty(size=(2*out_features, 1))) 64 | nn.init.xavier_uniform_(self.a.data, gain=1.414) 65 | 66 | self.leakyrelu = nn.LeakyReLU(self.alpha) 67 | 68 | def forward(self, h, adj): 69 | Wh = torch.mm(h, self.W) # h.shape: (N, in_features), Wh.shape: (N, out_features) 70 | e = self._prepare_attentional_mechanism_input(Wh) 71 | 72 | zero_vec = -9e15*torch.ones_like(e) 73 | attention = torch.where(adj > 0, e, zero_vec) 74 | attention = F.softmax(attention, dim=1) 75 | attention = F.dropout(attention, self.dropout, training=self.training) 76 | h_prime = torch.matmul(attention, Wh) 77 | 78 | if self.concat: 79 | return F.elu(h_prime) 80 | else: 81 | return h_prime 82 | 83 | def _prepare_attentional_mechanism_input(self, Wh): 84 | # Wh.shape (N, out_feature) 85 | # self.a.shape (2 * out_feature, 1) 86 | # Wh1&2.shape (N, 1) 87 | # e.shape (N, N) 88 | Wh1 = torch.matmul(Wh, self.a[:self.out_features, :]) 89 | Wh2 = torch.matmul(Wh, self.a[self.out_features:, :]) 90 | # broadcast add 91 | e = Wh1 + Wh2.T 92 | return self.leakyrelu(e) 93 | 94 | def __repr__(self): 95 | return self.__class__.__name__ + ' (' + str(self.in_features) + ' -> ' + str(self.out_features) + ')' 96 | 97 | 98 | # Graphsage layer 99 | class GraphSageConv(Module): 100 | """ 101 | Simple Graphsage layer 102 | """ 103 | 104 | def __init__(self, in_features, out_features, bias=False): 105 | super(GraphSageConv, self).__init__() 106 | 107 | self.proj = nn.Linear(in_features * 2, out_features, bias=bias) 108 | 109 | self.reset_parameters() 110 | 111 | # print("note: for dense graph in graphsage, require it normalized.") 112 | 113 | def reset_parameters(self): 114 | 115 | nn.init.normal_(self.proj.weight) 116 | 117 | if self.proj.bias is not None: 118 | nn.init.constant_(self.proj.bias, 0.) 119 | 120 | def forward(self, features, adj): 121 | """ 122 | Args: 123 | adj: can be sparse or dense matrix. 124 | """ 125 | 126 | # fuse info from neighbors. to be added: 127 | if not isinstance(adj, torch.sparse.Tensor): 128 | if len(adj.shape) == 3: 129 | neigh_feature = torch.bmm(adj, features) / ( 130 | adj.sum(dim=1).reshape((adj.shape[0], adj.shape[1], -1)) + 1) 131 | else: 132 | neigh_feature = torch.mm(adj, features) / (adj.sum(dim=1).reshape(adj.shape[0], -1) + 1) 133 | else: 134 | # print("spmm not implemented for batch training. Note!") 135 | 136 | neigh_feature = torch.spmm(adj, features) / (adj.to_dense().sum(dim=1).reshape(adj.shape[0], -1) + 1) 137 | 138 | # perform conv 139 | data = torch.cat([features, neigh_feature], dim=-1) 140 | combined = self.proj(data) 141 | 142 | return combined 143 | -------------------------------------------------------------------------------- /scripts/modules_model.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | import timm 3 | from transformers import DistilBertModel, DistilBertConfig 4 | from transformers import BertModel,BertConfig,BertTokenizer 5 | import torch.nn.functional as F 6 | from scripts.modules_graph import GraphConvolution,GraphAttentionLayer,GraphSageConv 7 | from utils import args 8 | from transformers import DistilBertTokenizer 9 | import cv2 10 | from matplotlib import pyplot as plt 11 | import torch 12 | 13 | 14 | class ImageEncoder(nn.Module): 15 | """ 16 | Encode images to a fixed size vector 17 | """ 18 | 19 | def __init__( 20 | self, transform,model_name=args.image_encoder_model, pretrained=args.pretrained, trainable=args.trainable 21 | ): 22 | super().__init__() 23 | # loading the pretrained model 24 | self.model = timm.create_model( 25 | model_name, pretrained, num_classes=0, global_pool="avg" 26 | ) 27 | 28 | for p in self.model.parameters(): 29 | p.requires_grad = trainable 30 | 31 | self.transforms = transform 32 | def forward(self, x): 33 | return self.model(x) 34 | 35 | def get_finetune_embed(self,img): 36 | 37 | self.model.train() 38 | image = plt.imread(img) 39 | image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) 40 | image = self.transforms(image=image)['image'] 41 | x = torch.tensor(image).permute(2, 0, 1).float().unsqueeze(0) 42 | 43 | return self.model(x) 44 | 45 | 46 | class TextEncoder(nn.Module): 47 | def __init__(self, model_name=args.text_encoder_model, pretrained=args.pretrained, trainable=args.trainable): 48 | super().__init__() 49 | if pretrained: 50 | if args.text_encoder_model == 'distilbert-base-uncased': 51 | self.model = DistilBertModel.from_pretrained(model_name) 52 | self.tokenizer = DistilBertTokenizer.from_pretrained(args.text_encoder_tokenizer) 53 | elif args.text_encoder_model == 'bert-base-uncased': 54 | self.model = BertModel.from_pretrained(model_name) 55 | self.tokenizer = BertTokenizer.from_pretrained(args.text_encoder_tokenizer) 56 | else: 57 | if args.text_encoder_model == 'distilbert-base-uncased': 58 | self.model = DistilBertModel(config=DistilBertConfig()) 59 | self.tokenizer = DistilBertTokenizer.from_pretrained(args.text_encoder_tokenizer) 60 | elif args.text_encoder_model == 'bert-base-uncased': 61 | self.model = BertModel(config=BertConfig()) 62 | self.tokenizer = BertTokenizer.from_pretrained(args.text_encoder_tokenizer) 63 | 64 | for p in self.model.parameters(): 65 | p.requires_grad = trainable 66 | 67 | # Using the CLS token hidden representation as the sentence's embedding 68 | self.target_token_idx = 0 69 | 70 | def forward(self, input_ids, attention_mask): 71 | output = self.model(input_ids=input_ids, attention_mask=attention_mask) 72 | last_hidden_state = output.last_hidden_state 73 | return last_hidden_state[:, self.target_token_idx, :] 74 | 75 | def get_finetune_embed(self,text): 76 | 77 | tokens = self.tokenizer(text, padding=True, truncation=True, max_length=args.max_length, return_tensors='pt').to(args.device) 78 | self.model.train() 79 | return self.model(**tokens).last_hidden_state[:,self.target_token_idx,:] 80 | 81 | 82 | class ProjectionHead(nn.Module): 83 | def __init__( 84 | self, 85 | embedding_dim, 86 | projection_dim=args.projection_dim, 87 | dropout=args.dropout 88 | ): 89 | super().__init__() 90 | self.projection = nn.Linear(embedding_dim, projection_dim) 91 | self.gelu = nn.GELU() 92 | self.fc = nn.Linear(projection_dim, projection_dim) 93 | self.dropout = nn.Dropout(dropout) 94 | self.layer_norm = nn.LayerNorm(projection_dim) 95 | 96 | def forward(self, x): 97 | projected = self.projection(x) 98 | x = self.gelu(projected) 99 | x = self.fc(x) 100 | x = self.dropout(x) 101 | x = x + projected 102 | x = self.layer_norm(x) 103 | return x 104 | 105 | class FeatureProjection(nn.Module): 106 | def __init__( 107 | self, feature, 108 | projection_dim = args.feat_projection_dim 109 | ): 110 | super().__init__() 111 | 112 | self.projection = {k: nn.Linear(v.shape[1], projection_dim).to(args.device) for k, v in feature.items()} 113 | self.gelu = nn.GELU() 114 | self.layer_norm = nn.LayerNorm(projection_dim) 115 | 116 | def forward(self, feature): 117 | project_feature = [] 118 | for k,v in feature.items(): 119 | projected = self.projection[k](v) 120 | x = self.gelu(projected) 121 | x = self.layer_norm(x) 122 | project_feature.append(x) 123 | 124 | if len(project_feature) >1: 125 | # return torch.cat([project_feature[0], project_feature[1]]) 126 | return torch.cat(project_feature,dim=0) 127 | else: 128 | return project_feature[0] 129 | # return torch.cat([project_feature[0],project_feature[1]]) 130 | 131 | class NodeEncoder(nn.Module): 132 | def __init__(self): 133 | super(NodeEncoder, self).__init__() 134 | 135 | self.dropout = args.node_dropout 136 | if args.node_encoder_model == 'gcn': 137 | self.gc1 = GraphConvolution(args.feat_projection_dim, args.hidden_dim) 138 | self.gc2 = GraphConvolution(args.hidden_dim, args.out_dim) 139 | 140 | elif args.node_encoder_model == 'gat': 141 | 142 | self.attentions = [GraphAttentionLayer(args.feat_projection_dim, args.hidden_dim, dropout=self.dropout, alpha=args.alpha, concat=True) for _ in 143 | range(args.nheads)] 144 | for i, attention in enumerate(self.attentions): 145 | self.add_module('attention_{}'.format(i), attention) 146 | 147 | self.out_att = GraphAttentionLayer(args.hidden_dim * args.nheads, args.out_dim, dropout=args.dropout, alpha=args.alpha, concat=False) 148 | elif args.node_encoder_model == 'sage': 149 | self.sage1 = GraphSageConv(args.feat_projection_dim, args.hidden_dim) 150 | self.sage2 = GraphSageConv(args.hidden_dim, args.out_dim) 151 | 152 | 153 | def forward(self, x, adj): 154 | if args.node_encoder_model == 'gcn': 155 | x1 = F.relu(self.gc1(x, adj)) 156 | x1 = F.dropout(x1, self.dropout, training=self.training) 157 | x2 = self.gc2(x1, adj) 158 | 159 | elif args.node_encoder_model == 'gat': 160 | x = F.dropout(x, self.dropout, training=self.training) 161 | x1 = torch.cat([att(x, adj) for att in self.attentions], dim=1) 162 | x1 = F.dropout(x1, self.dropout, training=self.training) 163 | x2 = F.elu(self.out_att(x1, adj)) 164 | 165 | elif args.node_encoder_model == 'sage': 166 | 167 | x1 = F.relu(self.sage1(x, adj)) 168 | x1 = F.dropout(x1, self.dropout, training=self.training) 169 | x2 = F.relu(self.sage2(x1, adj)) 170 | x2 = F.dropout(x2, self.dropout, training=self.training) 171 | 172 | return x2 173 | 174 | 175 | class NodeClassifier(nn.Module): 176 | def __init__(self): 177 | super(NodeClassifier, self).__init__() 178 | 179 | self.nodencoder = NodeEncoder() 180 | self.mlp = nn.Linear(args.out_dim, args.nclass) 181 | self.dropout = args.dropout 182 | 183 | self.reset_parameters() 184 | 185 | def reset_parameters(self): 186 | nn.init.normal_(self.mlp.weight, std=0.05) 187 | 188 | def forward(self, x, adj): 189 | x = F.relu(self.nodencoder(x, adj)) 190 | x = F.dropout(x, self.dropout, training=self.training) 191 | x = self.mlp(x[:,:]) 192 | 193 | return x 194 | 195 | 196 | class Classifier(nn.Module): 197 | def __init__(self,input_dim= args.out_dim,output_dim=args.nclass): 198 | super(Classifier, self).__init__() 199 | 200 | 201 | self.mlp = nn.Linear(input_dim, output_dim) 202 | self.dropout = args.dropout 203 | 204 | self.reset_parameters() 205 | 206 | def reset_parameters(self): 207 | nn.init.normal_(self.mlp.weight, std=0.05) 208 | 209 | def forward(self, x): 210 | x = F.relu(x) 211 | x = F.dropout(x, self.dropout, training=self.training) 212 | x = self.mlp(x[:,:]) 213 | 214 | return x 215 | 216 | 217 | class FineTuneProjection(nn.Module): 218 | def __init__( 219 | self, 220 | embedding_dim, 221 | projection_dim 222 | ): 223 | super().__init__() 224 | self.projection = nn.Linear(embedding_dim, projection_dim) 225 | self.gelu = nn.GELU() 226 | self.layer_norm = nn.LayerNorm(projection_dim) 227 | 228 | def forward(self, x): 229 | projected = self.projection(x) 230 | x = self.gelu(projected) 231 | x = self.layer_norm(x) 232 | return x 233 | -------------------------------------------------------------------------------- /scripts/moodules_graph.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from scripts.modules_layer import GraphConvolution,GraphSageConv,GraphAttentionLayer 5 | from utils.params import args 6 | 7 | class GCN(nn.Module): 8 | def __init__(self, nfeat, nhid, nclass, dropout,dataset,feature,projection_dim=200): 9 | super(GCN, self).__init__() 10 | 11 | self.gc1 = GraphConvolution(nfeat, nhid) 12 | self.gc2 = GraphConvolution(nhid, nhid) 13 | self.gc3 = GraphConvolution(nhid, nclass) 14 | self.dropout = dropout 15 | self.dataset = dataset 16 | if dataset == 'aminer': 17 | self.projection = FeatureProjection(feature,projection_dim).to(args.device) 18 | 19 | 20 | def forward(self, x, adj): 21 | if self.dataset == 'aminer': 22 | x = self.projection(x) 23 | x = F.relu(self.gc1(x, adj)) 24 | x = F.dropout(x, self.dropout, training=self.training) 25 | x =F.dropout(F.relu(self.gc2(x, adj))) 26 | x = self.gc3(x, adj) 27 | return F.log_softmax(x, dim=1) 28 | 29 | class SAGE(nn.Module): 30 | def __init__(self, nfeat, nhid, nclass, dropout,dataset,feature,projection_dim=200): 31 | super(SAGE, self).__init__() 32 | 33 | self.sage1 = GraphSageConv(nfeat, nhid) 34 | self.sage2 = GraphSageConv(nhid, nhid) 35 | self.mlp = nn.Linear(nhid, nclass) 36 | self.reset_parameters() 37 | 38 | self.dropout = dropout 39 | self.dataset = dataset 40 | if dataset == 'aminer': 41 | self.projection = FeatureProjection(feature,projection_dim).to(args.device) 42 | 43 | def reset_parameters(self): 44 | nn.init.normal_(self.mlp.weight,std=0.05) 45 | 46 | def forward(self, x, adj): 47 | if self.dataset == 'aminer': 48 | x = self.projection(x) 49 | x = F.relu(self.sage1(x, adj)) 50 | x = F.dropout(x, self.dropout, training=self.training) 51 | x = F.dropout(F.relu(self.sage2(x, adj))) 52 | x = self.mlp(x) 53 | 54 | return F.log_softmax(x, dim=1) 55 | 56 | class SAGE(nn.Module): 57 | def __init__(self, nfeat, nhid, nclass, dropout,dataset,feature,projection_dim=200): 58 | super(SAGE, self).__init__() 59 | 60 | self.sage1 = GraphSageConv(nfeat, nhid) 61 | self.sage2 = GraphSageConv(nhid, nhid) 62 | self.mlp = nn.Linear(nhid, nclass) 63 | self.reset_parameters() 64 | 65 | self.dropout = dropout 66 | self.dataset = dataset 67 | if dataset == 'aminer': 68 | self.projection = FeatureProjection(feature,projection_dim).to(args.device) 69 | 70 | def reset_parameters(self): 71 | nn.init.normal_(self.mlp.weight,std=0.05) 72 | 73 | def forward(self, x, adj): 74 | if self.dataset == 'aminer': 75 | x = self.projection(x) 76 | x = F.relu(self.sage1(x, adj)) 77 | x = F.dropout(x, self.dropout, training=self.training) 78 | x = F.dropout(F.relu(self.sage2(x, adj))) 79 | x = self.mlp(x) 80 | 81 | return F.log_softmax(x, dim=1) 82 | 83 | class GAT(nn.Module): 84 | def __init__(self, nfeat, nhid, nclass, dropout, dataset,feature,projection_dim=200,alpha=0.2, nheads=1): 85 | super(GAT, self).__init__() 86 | 87 | self.attentions = [GraphAttentionLayer(nfeat, nhid, dropout=dropout, alpha=alpha, concat=True) for _ in 88 | range(nheads)] 89 | for i, attention in enumerate(self.attentions): 90 | self.add_module('attention_{}'.format(i), attention) 91 | 92 | self.out_proj = nn.Linear(nhid * nheads, nhid) 93 | 94 | if args.dataset == 'aminer': 95 | self.projection = FeatureProjection(feature, projection_dim).to(args.device) 96 | 97 | self.dropout = dropout 98 | self.mlp = nn.Linear(nhid, nclass) 99 | self.dropout = dropout 100 | 101 | self.reset_parameters() 102 | 103 | self.dataset = dataset 104 | 105 | def reset_parameters(self): 106 | nn.init.normal_(self.mlp.weight, std=0.05) 107 | nn.init.normal_(self.out_proj.weight, std=0.05) 108 | 109 | def forward(self, x, adj): 110 | if self.dataset == 'aminer': 111 | x = self.projection(x) 112 | 113 | x = torch.cat([att(x, adj) for att in self.attentions], dim=1) 114 | x = F.dropout(x, self.dropout, training=self.training) 115 | x = F.elu(self.out_proj(x)) 116 | x = self.mlp(x) 117 | 118 | return F.log_softmax(x, dim=1) 119 | 120 | class FeatureProjection(nn.Module): 121 | def __init__( 122 | self, feature, 123 | projection_dim 124 | ): 125 | super().__init__() 126 | # device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 127 | self.projection = {k: nn.Linear(v.shape[1], projection_dim).to(args.device) for k, v in feature.items()} 128 | self.gelu = nn.GELU() 129 | self.layer_norm = nn.LayerNorm(projection_dim) 130 | 131 | def forward(self, feature): 132 | project_feature = [] 133 | for k,v in feature.items(): 134 | projected = self.projection[k](v) 135 | x = self.gelu(projected) 136 | x = self.layer_norm(x) 137 | 138 | project_feature.append(x) 139 | 140 | if len(project_feature) ==2: 141 | return torch.cat([project_feature[0],project_feature[1]]) 142 | elif len(project_feature) ==3: 143 | return torch.cat([project_feature[0],project_feature[1],project_feature[2]]) 144 | else: 145 | return torch.tensor(project_feature[0].clone().detach()) 146 | 147 | 148 | class GraphAttentionLayer(nn.Module): 149 | """ 150 | Simple GAT layer, similar to https://arxiv.org/abs/1710.10903 151 | """ 152 | 153 | def __init__(self, in_features, out_features, dropout, alpha, concat=True): 154 | super(GraphAttentionLayer, self).__init__() 155 | self.dropout = dropout 156 | self.in_features = in_features 157 | self.out_features = out_features 158 | self.alpha = alpha 159 | self.concat = concat 160 | 161 | self.W = nn.Parameter(torch.zeros(size=(in_features, out_features))) 162 | nn.init.xavier_uniform_(self.W.data, gain=1.414) 163 | self.a = nn.Parameter(torch.zeros(size=(2*out_features, 1))) 164 | nn.init.xavier_uniform_(self.a.data, gain=1.414) 165 | 166 | self.leakyrelu = nn.LeakyReLU(self.alpha) 167 | 168 | def forward(self, input, adj): 169 | if isinstance(adj, torch.sparse.FloatTensor): 170 | adj = adj.to_dense() 171 | 172 | h = torch.mm(input, self.W) 173 | N = h.size()[0] 174 | 175 | a_input = torch.cat([h.repeat(1, N).view(N * N, -1), h.repeat(N, 1)], dim=1).view(N, -1, 2 * self.out_features) 176 | e = self.leakyrelu(torch.matmul(a_input, self.a).squeeze(2)) 177 | 178 | zero_vec = -9e15*torch.ones_like(e) 179 | attention = torch.where(adj > 0, e, zero_vec) 180 | attention = F.softmax(attention, dim=1) 181 | attention = F.dropout(attention, self.dropout, training=self.training) 182 | h_prime = torch.matmul(attention, h) 183 | 184 | if self.concat: 185 | return F.elu(h_prime) 186 | else: 187 | return h_prime 188 | 189 | def __repr__(self): 190 | return self.__class__.__name__ + ' (' + str(self.in_features) + ' -> ' + str(self.out_features) + ')' 191 | 192 | 193 | 194 | -------------------------------------------------------------------------------- /utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .params import set_params,args 2 | -------------------------------------------------------------------------------- /utils/focalloss.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # encoding: utf-8 3 | 4 | from typing import Optional 5 | 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | from utils.getdata import encode_onehot 10 | import numpy as np 11 | import os 12 | os.environ['CUDA_VISIBLE_DEVICES'] = '1' 13 | 14 | 15 | # based on: 16 | # https://github.com/zhezh/focalloss/blob/master/focalloss.py 17 | 18 | class FocalLoss(nn.Module): 19 | r"""Criterion that computes Focal loss. 20 | 21 | According to [1], the Focal loss is computed as follows: 22 | 23 | .. math:: 24 | 25 | \text{FL}(p_t) = -\alpha_t (1 - p_t)^{\gamma} \, \text{log}(p_t) 26 | 27 | where: 28 | - :math:`p_t` is the model's estimated probability for each class. 29 | 30 | 31 | Arguments: 32 | alpha (float): Weighting factor :math:`\alpha \in [0, 1]`. 33 | gamma (float): Focusing parameter :math:`\gamma >= 0`. 34 | reduction (Optional[str]): Specifies the reduction to apply to the 35 | output: ‘none’ | ‘mean’ | ‘sum’. ‘none’: no reduction will be applied, 36 | ‘mean’: the sum of the output will be divided by the number of elements 37 | in the output, ‘sum’: the output will be summed. Default: ‘none’. 38 | 39 | Shape: 40 | - Input: :math:`(N, C, H, W)` where C = number of classes. 41 | - Target: :math:`(N, H, W)` where each value is 42 | :math:`0 ≤ targets[i] ≤ C−1`. 43 | 44 | Examples: 45 | >>> N = 5 # num_classes 46 | >>> loss = tgm.losses.FocalLoss(alpha=0.5, gamma=2.0, reduction='mean') 47 | >>> input = torch.randn(1, N, 3, 5, requires_grad=True) 48 | >>> target = torch.empty(1, 3, 5, dtype=torch.long).random_(N) 49 | >>> output = loss(input, target) 50 | >>> output.backward() 51 | 52 | References: 53 | [1] https://arxiv.org/abs/1708.02002 54 | """ 55 | 56 | def __init__(self, device, alpha: float, gamma: Optional[float] = 2.0, 57 | reduction: Optional[str] = 'none') -> None: 58 | super(FocalLoss, self).__init__() 59 | self.alpha: float = alpha 60 | self.gamma: Optional[float] = gamma 61 | self.reduction: Optional[str] = reduction 62 | self.eps: float = 1e-6 63 | self.device = device 64 | 65 | def forward( 66 | self, 67 | input: torch.Tensor, 68 | target: torch.Tensor) -> torch.Tensor: 69 | if not torch.is_tensor(input): 70 | raise TypeError("Input type is not a torch.Tensor. Got {}" 71 | .format(type(input))) 72 | 73 | # compute softmax over the classes axis 74 | input_soft = F.softmax(input, dim=1) + self.eps 75 | 76 | target_one_hot = torch.tensor(encode_onehot(np.array(target.cpu()))).to(self.device) 77 | # compute the actual focal loss 78 | weight = torch.pow(1. - input_soft, self.gamma) 79 | focal = -self.alpha * weight * torch.log(input_soft) 80 | loss_tmp = torch.sum(target_one_hot * focal, dim=1) 81 | 82 | loss = -1 83 | if self.reduction == 'none': 84 | loss = loss_tmp 85 | elif self.reduction == 'mean': 86 | loss = torch.mean(loss_tmp) 87 | elif self.reduction == 'sum': 88 | loss = torch.sum(loss_tmp) 89 | else: 90 | raise NotImplementedError("Invalid reduction mode: {}" 91 | .format(self.reduction)) 92 | return loss 93 | 94 | 95 | ###################### 96 | # functional interface 97 | ###################### 98 | 99 | 100 | def focal_loss( 101 | input: torch.Tensor, 102 | target: torch.Tensor, 103 | alpha: float, 104 | gamma: Optional[float] = 2.0, 105 | reduction: Optional[str] = 'none') -> torch.Tensor: 106 | r"""Function that computes Focal loss. 107 | 108 | See :class:`~torchgeometry.losses.FocalLoss` for details. 109 | """ 110 | return FocalLoss(alpha, gamma, reduction)(input, target) 111 | 112 | 113 | def one_hot(labels: torch.Tensor, 114 | num_classes: int, 115 | device: Optional[torch.device] = None, 116 | dtype: Optional[torch.dtype] = None, 117 | eps: Optional[float] = 1e-6) -> torch.Tensor: 118 | r"""Converts an integer label 2D tensor to a one-hot 3D tensor. 119 | 120 | Args: 121 | labels (torch.Tensor) : tensor with labels of shape :math:`(N, H, W)`, 122 | where N is batch siz. Each value is an integer 123 | representing correct classification. 124 | num_classes (int): number of classes in labels. 125 | device (Optional[torch.device]): the desired device of returned tensor. 126 | Default: if None, uses the current device for the default tensor type 127 | (see torch.set_default_tensor_type()). device will be the CPU for CPU 128 | tensor types and the current CUDA device for CUDA tensor types. 129 | dtype (Optional[torch.dtype]): the desired data type of returned 130 | tensor. Default: if None, infers data type from values. 131 | 132 | Returns: 133 | torch.Tensor: the labels in one hot tensor. 134 | 135 | Examples:: 136 | >>> labels = torch.LongTensor([[[0, 1], [2, 0]]]) 137 | >>> tgm.losses.one_hot(labels, num_classes=3) 138 | tensor([[[[1., 0.], 139 | [0., 1.]], 140 | [[0., 1.], 141 | [0., 0.]], 142 | [[0., 0.], 143 | [1., 0.]]]] 144 | """ 145 | if not torch.is_tensor(labels): 146 | raise TypeError("Input labels type is not a torch.Tensor. Got {}" 147 | .format(type(labels))) 148 | if not len(labels.shape) == 3: 149 | raise ValueError("Invalid depth shape, we expect BxHxW. Got: {}" 150 | .format(labels.shape)) 151 | if not labels.dtype == torch.int64: 152 | raise ValueError( 153 | "labels must be of the same dtype torch.int64. Got: {}" .format( 154 | labels.dtype)) 155 | if num_classes < 1: 156 | raise ValueError("The number of classes must be bigger than one." 157 | " Got: {}".format(num_classes)) 158 | batch_size, height, width = labels.shape 159 | one_hot = torch.zeros(batch_size, num_classes, height, width, 160 | device=device, dtype=dtype) 161 | return one_hot.scatter_(1, labels.unsqueeze(1), 1.0) + eps 162 | -------------------------------------------------------------------------------- /utils/getdata.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | from matplotlib import pyplot as plt 3 | import albumentations as A 4 | import numpy as np 5 | import scipy.sparse as sp 6 | import torch 7 | import random 8 | from utils.params import args 9 | 10 | 11 | class NodeImageDataset(): 12 | def __init__(self, entity_id,image,label, transforms): 13 | 14 | 15 | self.image_filenames = list(image) 16 | self.entity_id = entity_id 17 | self.transforms = transforms 18 | self.label = label 19 | 20 | def __getitem__(self, idx): 21 | 22 | item = {} 23 | # print(idx) 24 | 25 | try: 26 | image = plt.imread(self.image_filenames[idx]) 27 | image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) 28 | image = self.transforms(image=image)['image'] 29 | except: 30 | image = np.random.randn(224,224,3) 31 | # print("error happened") 32 | item['image'] = torch.tensor(image).permute(2, 0, 1).float() 33 | item['entity'] = self.entity_id[idx] 34 | item['label'] = self.label[idx] 35 | 36 | return item 37 | 38 | def __len__(self): 39 | return len(self.image_filenames) 40 | 41 | 42 | class NodeTextDataset(): 43 | def __init__(self, entity_id,text,label,tokenizer): 44 | 45 | self.text = list(text) 46 | 47 | self.encoded_text = tokenizer( 48 | self.text, padding=True, truncation=True, max_length=args.max_length, return_tensors='pt' 49 | ) 50 | 51 | self.entity_id = entity_id 52 | self.label = label 53 | 54 | def __getitem__(self, idx): 55 | item = { 56 | key: torch.tensor(values[idx]) 57 | for key, values in self.encoded_text.items() 58 | } 59 | item['text'] = self.text[idx] 60 | item['entity'] = self.entity_id[idx] 61 | item['label'] = self.label[idx] 62 | 63 | return item 64 | 65 | 66 | def __len__(self): 67 | return len(self.text) 68 | 69 | 70 | # get transformation for image data 71 | def get_transforms(mode="train"): 72 | if mode == "train": 73 | return A.Compose( 74 | [ 75 | A.Resize(args.image_size, args.image_size, always_apply=True), 76 | A.Normalize(max_pixel_value=255.0, always_apply=True), 77 | ] 78 | ) 79 | else: 80 | return A.Compose( 81 | [ 82 | A.Resize(args.image_size, args.image_size, always_apply=True), 83 | A.Normalize(max_pixel_value=255.0, always_apply=True), 84 | ] 85 | ) 86 | 87 | def split_imbalance(labels,train_ratio,val_ratio,test_ratio,imbalance_ratio): 88 | 89 | num_classes = len(set(labels.tolist())) 90 | c_idxs = [] # class-wise index 91 | train_idx = [] 92 | val_idx = [] 93 | test_idx = [] 94 | c_num_mat = np.zeros((num_classes,3)).astype(int) 95 | label_max = int(max(labels.tolist())+1) 96 | minority_index = [item for item in range(label_max) if labels.tolist().count(item) < len(labels.tolist())/num_classes] 97 | 98 | for i in range(num_classes): 99 | c_idx = (labels==i).nonzero()[:,-1].tolist() 100 | c_num = len(c_idx) 101 | if num_classes > 2: 102 | if i in minority_index: c_num = int(len(labels.tolist()) / num_classes * imbalance_ratio) 103 | else: 104 | if i in minority_index: c_num = int((len(labels.tolist())- labels.tolist().count(i)) * imbalance_ratio) 105 | 106 | print('The number of class {:d}: {:d}'.format(i,c_num)) 107 | random.shuffle(c_idx) 108 | c_idxs.append(c_idx) 109 | 110 | if c_num <4: 111 | if c_num < 3: 112 | print("too small class type") 113 | c_num_mat[i,0] = 1 114 | c_num_mat[i,1] = 1 115 | c_num_mat[i,2] = 1 116 | else: 117 | c_num_mat[i,0] = int(c_num/10 *train_ratio) 118 | c_num_mat[i,1] = int(c_num/10 * val_ratio) 119 | c_num_mat[i,2] = int(c_num/10 * test_ratio) 120 | 121 | 122 | train_idx = train_idx + c_idx[:c_num_mat[i,0]] 123 | 124 | val_idx = val_idx + c_idx[c_num_mat[i,0]:c_num_mat[i,0]+c_num_mat[i,1]] 125 | test_idx = test_idx + c_idx[c_num_mat[i,0]+c_num_mat[i,1]:c_num_mat[i,0]+c_num_mat[i,1]+c_num_mat[i,2]] 126 | 127 | random.shuffle(train_idx) 128 | 129 | 130 | return train_idx, val_idx, test_idx, c_num_mat 131 | 132 | def split_balance(labels,train_ratio,val_ratio,test_ratio): 133 | 134 | num_classes = len(set(labels.tolist())) 135 | c_idxs = [] 136 | train_idx = [] 137 | val_idx = [] 138 | test_idx = [] 139 | c_num_mat = np.zeros((num_classes,3)).astype(int) 140 | count_0, count_1 = labels.tolist().count(0), labels.tolist().count(1) 141 | count_min = min(count_0,count_1) 142 | 143 | for i in range(num_classes): 144 | c_idx = (labels==i).nonzero()[:,-1].tolist() 145 | c_num = count_min 146 | 147 | print('The number of class {:d}: {:d}'.format(i,c_num)) 148 | 149 | random.shuffle(c_idx) 150 | c_idxs.append(c_idx) 151 | 152 | if c_num <4: 153 | if c_num < 3: 154 | print("too small class type") 155 | c_num_mat[i,0] = 1 156 | c_num_mat[i,1] = 1 157 | c_num_mat[i,2] = 1 158 | else: 159 | c_num_mat[i,0] = int(c_num/10 *train_ratio) 160 | c_num_mat[i,1] = int(c_num/10 * val_ratio) 161 | c_num_mat[i,2] = int(c_num/10 * test_ratio) 162 | 163 | 164 | train_idx = train_idx + c_idx[:c_num_mat[i,0]] 165 | 166 | val_idx = val_idx + c_idx[c_num_mat[i,0]:c_num_mat[i,0]+c_num_mat[i,1]] 167 | test_idx = test_idx + c_idx[c_num_mat[i,0]+c_num_mat[i,1]:c_num_mat[i,0]+c_num_mat[i,1]+c_num_mat[i,2]] 168 | 169 | 170 | return train_idx, val_idx, test_idx, c_num_mat 171 | 172 | def split_genuine(labels,train_ratio,val_ratio,test_ratio): 173 | 174 | num_classes = len(set(labels.tolist())) 175 | c_idxs = [] # class-wise index 176 | train_idx = [] 177 | val_idx = [] 178 | test_idx = [] 179 | c_num_mat = np.zeros((num_classes,3)).astype(int) 180 | 181 | for i in range(num_classes): 182 | c_idx = (labels==i).nonzero()[:,-1].tolist() 183 | c_num = len(c_idx) 184 | print('The number of class {:d}: {:d}'.format(i,c_num)) 185 | random.shuffle(c_idx) 186 | c_idxs.append(c_idx) 187 | 188 | if c_num <4: 189 | if c_num < 3: 190 | print("too small class type") 191 | c_num_mat[i,0] = 1 192 | c_num_mat[i,1] = 1 193 | c_num_mat[i,2] = 1 194 | else: 195 | c_num_mat[i,0] = int(c_num/10 *train_ratio) 196 | c_num_mat[i,1] = int(c_num/10 * val_ratio) 197 | c_num_mat[i,2] = int(c_num/10 * test_ratio) 198 | 199 | 200 | train_idx = train_idx + c_idx[:c_num_mat[i,0]] 201 | 202 | val_idx = val_idx + c_idx[c_num_mat[i,0]:c_num_mat[i,0]+c_num_mat[i,1]] 203 | test_idx = test_idx + c_idx[c_num_mat[i,0]+c_num_mat[i,1]:c_num_mat[i,0]+c_num_mat[i,1]+c_num_mat[i,2]] 204 | 205 | 206 | return train_idx, val_idx, test_idx, c_num_mat 207 | 208 | def load_data_for_pretrain(): 209 | 210 | if args.dataset == 'instagram': 211 | embed_dir = args.feature_path 212 | relation_dir = args.relation_path 213 | pos_dir = args.pos_path 214 | 215 | idx_features_labels = np.round(np.genfromtxt(embed_dir, 216 | dtype=np.float, delimiter=' ', invalid_raise=True),4) 217 | 218 | features = idx_features_labels[:, 770:] 219 | features = sp.csr_matrix(features, dtype=np.float32) 220 | total_num = features.shape[0] 221 | features = {'feature': torch.FloatTensor(np.array(features.todense())).to(args.device)} 222 | number = 8225 223 | label = encode_onehot(idx_features_labels[:number, 1]) 224 | 225 | # build graph 226 | idx = np.array(idx_features_labels[:, 0], dtype=np.float) 227 | idx_map = {j: i for i, j in enumerate(idx)} 228 | 229 | 230 | elif args.dataset == 'github': 231 | 232 | embed_dir = args.feature_path 233 | 234 | relation_dir = args.relation_path 235 | pos_dir = args.pos_path 236 | 237 | idx_features_labels = np.genfromtxt(embed_dir, 238 | dtype=np.dtype(str), delimiter=',', invalid_raise=True) 239 | 240 | features = sp.csr_matrix(idx_features_labels[:, :], dtype=np.float32) 241 | 242 | label_data = idx_features_labels[:,0] 243 | label = encode_onehot(label_data) 244 | 245 | # build graph 246 | idx = np.array(list(range(idx_features_labels.shape[0])), dtype=np.int32) 247 | idx_map = {j: i for i, j in enumerate(idx)} 248 | 249 | total_num = features.shape[0] 250 | 251 | elif args.dataset == 'yelp': 252 | 253 | fea_dir = args.feature_path 254 | relation_dir = args.relation_path 255 | idx_features_labels = np.genfromtxt(fea_dir, dtype=np.dtype(str), delimiter=' ', invalid_raise=True) 256 | features = idx_features_labels[:, 2:] 257 | pos_dir = args.pos_path 258 | if args.yelp_concate: 259 | embed_dir = args.embed_path 260 | idx_embeds_labels = np.genfromtxt(embed_dir, 261 | dtype=np.dtype(str), delimiter=' ', invalid_raise=True) 262 | embeds = idx_embeds_labels[:, 2:] 263 | 264 | features = np.concatenate((features, embeds), axis=1) 265 | 266 | features = sp.csr_matrix(features, dtype=np.float32) 267 | total_num = features.shape[0] 268 | features = normalize(features) 269 | 270 | number = 67395 271 | label = torch.tensor(encode_onehot(idx_features_labels[:number, 1])) 272 | 273 | # build graph 274 | idx = np.array([i for i in range(features.shape[0])]) 275 | idx_map = {j: i for i, j in enumerate(idx)} 276 | 277 | features = {'feature': torch.FloatTensor(np.array(features.todense())).to(args.device)} 278 | 279 | elif args.dataset == 'aminer': 280 | 281 | 282 | fea_dir = args.feature_path 283 | author_fea_dir = args.author_feature_path 284 | 285 | relation_dir = args.relation_path 286 | idx_features_labels = np.genfromtxt(fea_dir, 287 | dtype=np.dtype(str), delimiter=',', invalid_raise=True, skip_header=True) 288 | paper_features = idx_features_labels[:, 2:] 289 | author_features = np.genfromtxt(author_fea_dir, dtype=np.dtype(str), delimiter=',', invalid_raise=True, 290 | skip_header=True)[:, 2:] 291 | 292 | number = 18089 293 | 294 | paper_features = sp.csr_matrix(paper_features, dtype=np.float32) 295 | author_features = sp.csr_matrix(author_features, dtype=np.float32) 296 | paper_features = normalize(paper_features) 297 | paper_features = torch.FloatTensor(np.array(paper_features.todense())) 298 | author_features = torch.FloatTensor(np.array(author_features.todense())) 299 | features = {'paper_feature': paper_features.to(args.device), 'author_features': author_features.to(args.device)} 300 | label = encode_onehot(idx_features_labels[:number, 1].astype('float')) 301 | 302 | # build graph 303 | idx = np.array([i for i in range(paper_features.shape[0] + author_features.shape[0])]) 304 | total_num = paper_features.shape[0] + author_features.shape[0] 305 | idx_map = {j: i for i, j in enumerate(idx)} 306 | 307 | pos_dir = args.pos_path 308 | 309 | edges_unordered = np.genfromtxt(relation_dir, 310 | dtype=np.int32)[:, :-1] 311 | edges = np.array(list(map(idx_map.get, edges_unordered.flatten())), 312 | dtype=np.int32).reshape(edges_unordered.shape) 313 | adj = sp.coo_matrix((np.ones(edges.shape[0]), (edges[:, 0], edges[:, 1])), 314 | shape=(total_num,total_num), 315 | dtype=np.float32) 316 | 317 | # build symmetric adjacency matrix 318 | adj = adj + adj.T.multiply(adj.T > adj) - adj.multiply(adj.T > adj) 319 | 320 | # features = normalize(features) 321 | adj = normalize(adj + sp.eye(adj.shape[0])) 322 | 323 | pos_unordered = np.genfromtxt(pos_dir, 324 | dtype=np.int32) 325 | pos = np.array(list(map(idx_map.get, pos_unordered.flatten())), 326 | dtype=np.int32).reshape(pos_unordered.shape) 327 | pos = sp.coo_matrix((np.ones(pos.shape[0]), (pos[:, 0], pos[:, 1])), 328 | shape=(number,number), 329 | dtype=np.int32) 330 | # build symmetric adjacency matrix 331 | pos = pos + pos.T.multiply(pos.T > pos) - pos.multiply(pos.T > pos) 332 | pos = pos + sp.eye(pos.shape[0]) 333 | 334 | 335 | labels = torch.LongTensor(np.where(label)[1]).to(args.device) 336 | adj = sparse_mx_to_torch_sparse_tensor(adj).to(args.device) 337 | pos = sparse_mx_to_torch_sparse_tensor(pos).to(args.device) 338 | 339 | return adj, features, labels, pos 340 | 341 | def load_data_for_finetune(finetuned_feature_dir): 342 | 343 | if args.dataset == 'instagram': 344 | embed_dir = args.feature_path 345 | relation_dir = args.relation_path 346 | idx_features_labels = np.genfromtxt(embed_dir, 347 | dtype=np.dtype(str), delimiter=' ', invalid_raise=True) 348 | features = sp.csr_matrix(idx_features_labels[:, 2:], dtype=np.float32) 349 | number = 8225 350 | label = encode_onehot(idx_features_labels[:number, 1]) 351 | 352 | # build graph 353 | idx = np.array(idx_features_labels[:, 0], dtype=np.float) 354 | idx_map = {j: i for i, j in enumerate(idx)} 355 | 356 | elif args.dataset == 'github': 357 | 358 | embed_dir = args.feature_path 359 | 360 | relation_dir = args.relation_path 361 | 362 | idx_features_labels = np.genfromtxt(embed_dir, 363 | dtype=np.dtype(str), delimiter=',', invalid_raise=True) 364 | 365 | features = sp.csr_matrix(idx_features_labels[:, :], dtype=np.float32) 366 | 367 | label_data = idx_features_labels[:,0] 368 | 369 | label = encode_onehot(label_data) 370 | 371 | # build graph 372 | idx = np.array(list(range(idx_features_labels.shape[0])), dtype=np.int32) 373 | # idx = np.array(idx_features_labels[:, 0], dtype=np.float) 374 | idx_map = {j: i for i, j in enumerate(idx)} 375 | 376 | elif args.dataset == 'yelp': 377 | 378 | fea_dir = args.feature_path 379 | relation_dir = args.relation_path 380 | number = 67395 381 | 382 | idx_features_labels = np.genfromtxt(fea_dir, dtype=np.dtype(str), delimiter=' ', invalid_raise=True) 383 | features = idx_features_labels[:, 2:] 384 | features = normalize(features) 385 | features = sp.csr_matrix(features, dtype=np.float32) 386 | 387 | total_num = features.shape[0] 388 | label_data = idx_features_labels[:number,1] 389 | label = torch.tensor(encode_onehot(label_data)) 390 | 391 | # build graph 392 | idx = np.array([i for i in range(features.shape[0])]) 393 | idx_map = {j: i for i, j in enumerate(idx)} 394 | 395 | features = {'feature': torch.FloatTensor(np.array(features.todense())).to(args.device)} 396 | 397 | elif args.dataset == 'aminer': 398 | 399 | fea_dir = args.feature_path 400 | 401 | author_fea_dir = args.author_feature_path 402 | 403 | relation_dir = args.relation_path 404 | idx_features_labels = np.genfromtxt(fea_dir, 405 | dtype=np.dtype(str), delimiter=',', invalid_raise=True, skip_header=True) 406 | features = torch.from_numpy(np.genfromtxt(finetuned_feature_dir, delimiter=' ')).float() 407 | features = sp.csr_matrix(features, dtype=np.float32) 408 | features = normalize(features) 409 | features = torch.FloatTensor(np.array(features.todense())) 410 | features = {'features': features.to(args.device)} 411 | 412 | paper_features = idx_features_labels[:, 2:] 413 | author_features = np.genfromtxt(author_fea_dir, dtype=np.dtype(str), delimiter=',', invalid_raise=True, 414 | skip_header=True)[:, 2:] 415 | number = 18089 416 | 417 | # paper_features = sp.csr_matrix(paper_features, dtype=np.float32) 418 | # author_features = sp.csr_matrix(author_features, dtype=np.float32) 419 | # paper_features = normalize(paper_features) 420 | # paper_features = torch.FloatTensor(np.array(paper_features.todense())) 421 | # author_features = torch.FloatTensor(np.array(author_features.todense())) 422 | # features = {'paper_feature': paper_features.to(args.device), 'author_features': author_features.to(args.device)} 423 | label_data = idx_features_labels[:number, 1].astype('float') 424 | label = encode_onehot(label_data) 425 | 426 | 427 | 428 | # build graph 429 | # idx = np.array([i for i in range(paper_features.shape[0] + author_features.shape[0])]) 430 | # total_num = paper_features.shape[0] + author_features.shape[0] 431 | 432 | idx = np.array([i for i in range(paper_features.shape[0] + author_features.shape[0])]) 433 | total_num = paper_features.shape[0] + author_features.shape[0] 434 | idx_map = {j: i for i, j in enumerate(idx)} 435 | 436 | 437 | edges_unordered = np.genfromtxt(relation_dir, 438 | dtype=np.int32)[:, :-1] 439 | 440 | edges = np.array(list(map(idx_map.get, edges_unordered.flatten())), 441 | dtype=np.int32).reshape(edges_unordered.shape) 442 | 443 | adj = sp.coo_matrix((np.ones(edges.shape[0]), (edges[:, 0], edges[:, 1])), 444 | shape=(total_num, total_num), 445 | dtype=np.float32) 446 | 447 | # build symmetric adjacency matrix 448 | adj = adj + adj.T.multiply(adj.T > adj) - adj.multiply(adj.T > adj) 449 | 450 | # features = normalize(features) 451 | adj = normalize(adj + sp.eye(adj.shape[0])) 452 | 453 | if args.dataset == 'yelp': 454 | idx_train, idx_val, idx_test, c_num_mat = split_genuine(torch.tensor(label_data), train_ratio=4, val_ratio=2, 455 | test_ratio=4) 456 | else: 457 | idx_train, idx_val, idx_test, c_num_mat = split_genuine(torch.tensor(label_data),train_ratio=7,val_ratio=1,test_ratio=2) 458 | 459 | idx_train = torch.LongTensor(idx_train) 460 | idx_val = torch.LongTensor(idx_val) 461 | idx_test = torch.LongTensor(idx_test) 462 | 463 | labels = torch.LongTensor(np.where(label)[1]) 464 | adj = sparse_mx_to_torch_sparse_tensor(adj) 465 | 466 | return adj, features, labels, idx_train, idx_val, idx_test 467 | 468 | 469 | def load_data_for_finetune_imbalance(finetuned_feature_dir): 470 | 471 | if args.dataset == 'instagram': 472 | embed_dir = args.feature_path 473 | relation_dir = args.relation_path 474 | idx_features_labels = np.genfromtxt(embed_dir, 475 | dtype=np.dtype(str), delimiter=' ', invalid_raise=True) 476 | features = sp.csr_matrix(idx_features_labels[:, 2:], dtype=np.float32) 477 | number = 8225 478 | label = encode_onehot(idx_features_labels[:number, 1]) 479 | 480 | # build graph 481 | idx = np.array(idx_features_labels[:, 0], dtype=np.float) 482 | idx_map = {j: i for i, j in enumerate(idx)} 483 | 484 | elif args.dataset == 'github': 485 | 486 | embed_dir = args.feature_path 487 | 488 | relation_dir = args.relation_path 489 | 490 | idx_features_labels = np.genfromtxt(embed_dir, 491 | dtype=np.dtype(str), delimiter=',', invalid_raise=True) 492 | 493 | features = sp.csr_matrix(idx_features_labels[:, :], dtype=np.float32) 494 | 495 | label_data = idx_features_labels[:,0] 496 | 497 | label = encode_onehot(label_data) 498 | 499 | # build graph 500 | idx = np.array(list(range(idx_features_labels.shape[0])), dtype=np.int32) 501 | # idx = np.array(idx_features_labels[:, 0], dtype=np.float) 502 | idx_map = {j: i for i, j in enumerate(idx)} 503 | 504 | elif args.dataset == 'yelp': 505 | 506 | fea_dir = args.feature_path 507 | relation_dir = args.relation_path 508 | number = 67395 509 | 510 | idx_features_labels = np.genfromtxt(fea_dir, dtype=np.dtype(str), delimiter=' ', invalid_raise=True) 511 | features = idx_features_labels[:, 2:] 512 | features = normalize(features) 513 | features = sp.csr_matrix(features, dtype=np.float32) 514 | 515 | total_num = features.shape[0] 516 | label_data = idx_features_labels[:number,1] 517 | label = torch.tensor(encode_onehot(label_data)) 518 | 519 | # build graph 520 | idx = np.array([i for i in range(features.shape[0])]) 521 | idx_map = {j: i for i, j in enumerate(idx)} 522 | 523 | features = {'feature': torch.FloatTensor(np.array(features.todense())).to(args.device)} 524 | 525 | elif args.dataset == 'aminer': 526 | 527 | fea_dir = args.feature_path 528 | 529 | author_fea_dir = args.author_feature_path 530 | 531 | relation_dir = args.relation_path 532 | idx_features_labels = np.genfromtxt(fea_dir, 533 | dtype=np.dtype(str), delimiter=',', invalid_raise=True, skip_header=True) 534 | features = torch.from_numpy(np.genfromtxt(finetuned_feature_dir, delimiter=' ')).float() 535 | features = sp.csr_matrix(features, dtype=np.float32) 536 | features = normalize(features) 537 | features = torch.FloatTensor(np.array(features.todense())) 538 | features = {'features': features.to(args.device)} 539 | 540 | paper_features = idx_features_labels[:, 2:] 541 | author_features = np.genfromtxt(author_fea_dir, dtype=np.dtype(str), delimiter=',', invalid_raise=True, 542 | skip_header=True)[:, 2:] 543 | number = 18089 544 | 545 | # paper_features = sp.csr_matrix(paper_features, dtype=np.float32) 546 | # author_features = sp.csr_matrix(author_features, dtype=np.float32) 547 | # paper_features = normalize(paper_features) 548 | # paper_features = torch.FloatTensor(np.array(paper_features.todense())) 549 | # author_features = torch.FloatTensor(np.array(author_features.todense())) 550 | # features = {'paper_feature': paper_features.to(args.device), 'author_features': author_features.to(args.device)} 551 | label_data = idx_features_labels[:number, 1].astype('float') 552 | label = encode_onehot(label_data) 553 | 554 | 555 | 556 | # build graph 557 | # idx = np.array([i for i in range(paper_features.shape[0] + author_features.shape[0])]) 558 | # total_num = paper_features.shape[0] + author_features.shape[0] 559 | 560 | idx = np.array([i for i in range(paper_features.shape[0] + author_features.shape[0])]) 561 | total_num = paper_features.shape[0] + author_features.shape[0] 562 | idx_map = {j: i for i, j in enumerate(idx)} 563 | 564 | 565 | edges_unordered = np.genfromtxt(relation_dir, 566 | dtype=np.int32)[:, :-1] 567 | 568 | edges = np.array(list(map(idx_map.get, edges_unordered.flatten())), 569 | dtype=np.int32).reshape(edges_unordered.shape) 570 | 571 | adj = sp.coo_matrix((np.ones(edges.shape[0]), (edges[:, 0], edges[:, 1])), 572 | shape=(total_num, total_num), 573 | dtype=np.float32) 574 | 575 | # build symmetric adjacency matrix 576 | adj = adj + adj.T.multiply(adj.T > adj) - adj.multiply(adj.T > adj) 577 | 578 | # features = normalize(features) 579 | adj = normalize(adj + sp.eye(adj.shape[0])) 580 | 581 | if args.dataset == 'yelp': 582 | # idx_train, idx_val, idx_test, c_num_mat = split_genuine(torch.tensor(label_data),train_ratio=7,val_ratio=1,test_ratio=2) 583 | idx_train, idx_val, idx_test, c_num_mat = split_imbalance(torch.tensor(label_data), train_ratio=4, val_ratio=2, 584 | test_ratio=2, imbalance_ratio=args.imbalance_ratio) 585 | else: 586 | idx_train, idx_val, idx_test, c_num_mat = split_imbalance(torch.tensor(label_data), train_ratio=7, val_ratio=1, 587 | test_ratio=2, imbalance_ratio=args.imbalance_ratio) 588 | 589 | idx_train = torch.LongTensor(idx_train) 590 | idx_val = torch.LongTensor(idx_val) 591 | idx_test = torch.LongTensor(idx_test) 592 | 593 | labels = torch.LongTensor(np.where(label)[1]) 594 | adj = sparse_mx_to_torch_sparse_tensor(adj) 595 | 596 | return adj, features, labels, idx_train, idx_val, idx_test 597 | 598 | def load_data_for_finetune_balance(finetuned_feature_dir): 599 | 600 | if args.dataset == 'instagram': 601 | embed_dir = args.feature_path 602 | relation_dir = args.relation_path 603 | idx_features_labels = np.genfromtxt(embed_dir, 604 | dtype=np.dtype(str), delimiter=' ', invalid_raise=True) 605 | features = sp.csr_matrix(idx_features_labels[:, 2:], dtype=np.float32) 606 | number = 8225 607 | label = encode_onehot(idx_features_labels[:number, 1]) 608 | 609 | # build graph 610 | idx = np.array(idx_features_labels[:, 0], dtype=np.float) 611 | idx_map = {j: i for i, j in enumerate(idx)} 612 | 613 | elif args.dataset == 'github': 614 | 615 | embed_dir = args.feature_path 616 | 617 | relation_dir = args.relation_path 618 | 619 | idx_features_labels = np.genfromtxt(embed_dir, 620 | dtype=np.dtype(str), delimiter=',', invalid_raise=True) 621 | 622 | features = sp.csr_matrix(idx_features_labels[:, :], dtype=np.float32) 623 | 624 | label_data = idx_features_labels[:,0] 625 | 626 | label = encode_onehot(label_data) 627 | 628 | # build graph 629 | idx = np.array(list(range(idx_features_labels.shape[0])), dtype=np.int32) 630 | # idx = np.array(idx_features_labels[:, 0], dtype=np.float) 631 | idx_map = {j: i for i, j in enumerate(idx)} 632 | 633 | elif args.dataset == 'yelp': 634 | 635 | fea_dir = args.feature_path 636 | relation_dir = args.relation_path 637 | number = 67395 638 | 639 | idx_features_labels = np.genfromtxt(fea_dir, dtype=np.dtype(str), delimiter=' ', invalid_raise=True) 640 | features = idx_features_labels[:, 2:] 641 | features = normalize(features) 642 | features = sp.csr_matrix(features, dtype=np.float32) 643 | 644 | total_num = features.shape[0] 645 | label_data = idx_features_labels[:number,1] 646 | label = torch.tensor(encode_onehot(label_data)) 647 | 648 | # build graph 649 | idx = np.array([i for i in range(features.shape[0])]) 650 | idx_map = {j: i for i, j in enumerate(idx)} 651 | 652 | features = {'feature': torch.FloatTensor(np.array(features.todense())).to(args.device)} 653 | 654 | elif args.dataset == 'aminer': 655 | 656 | fea_dir = args.feature_path 657 | 658 | author_fea_dir = args.author_feature_path 659 | 660 | relation_dir = args.relation_path 661 | idx_features_labels = np.genfromtxt(fea_dir, 662 | dtype=np.dtype(str), delimiter=',', invalid_raise=True, skip_header=True) 663 | features = torch.from_numpy(np.genfromtxt(finetuned_feature_dir, delimiter=' ')).float() 664 | features = sp.csr_matrix(features, dtype=np.float32) 665 | features = normalize(features) 666 | features = torch.FloatTensor(np.array(features.todense())) 667 | features = {'features': features.to(args.device)} 668 | 669 | paper_features = idx_features_labels[:, 2:] 670 | author_features = np.genfromtxt(author_fea_dir, dtype=np.dtype(str), delimiter=',', invalid_raise=True, 671 | skip_header=True)[:, 2:] 672 | number = 18089 673 | 674 | # paper_features = sp.csr_matrix(paper_features, dtype=np.float32) 675 | # author_features = sp.csr_matrix(author_features, dtype=np.float32) 676 | # paper_features = normalize(paper_features) 677 | # paper_features = torch.FloatTensor(np.array(paper_features.todense())) 678 | # author_features = torch.FloatTensor(np.array(author_features.todense())) 679 | # features = {'paper_feature': paper_features.to(args.device), 'author_features': author_features.to(args.device)} 680 | label_data = idx_features_labels[:number, 1].astype('float') 681 | label = encode_onehot(label_data) 682 | 683 | 684 | 685 | # build graph 686 | # idx = np.array([i for i in range(paper_features.shape[0] + author_features.shape[0])]) 687 | # total_num = paper_features.shape[0] + author_features.shape[0] 688 | 689 | idx = np.array([i for i in range(paper_features.shape[0] + author_features.shape[0])]) 690 | total_num = paper_features.shape[0] + author_features.shape[0] 691 | idx_map = {j: i for i, j in enumerate(idx)} 692 | 693 | 694 | edges_unordered = np.genfromtxt(relation_dir, 695 | dtype=np.int32)[:, :-1] 696 | 697 | edges = np.array(list(map(idx_map.get, edges_unordered.flatten())), 698 | dtype=np.int32).reshape(edges_unordered.shape) 699 | 700 | adj = sp.coo_matrix((np.ones(edges.shape[0]), (edges[:, 0], edges[:, 1])), 701 | shape=(total_num, total_num), 702 | dtype=np.float32) 703 | 704 | # build symmetric adjacency matrix 705 | adj = adj + adj.T.multiply(adj.T > adj) - adj.multiply(adj.T > adj) 706 | 707 | # features = normalize(features) 708 | adj = normalize(adj + sp.eye(adj.shape[0])) 709 | 710 | if args.dataset == 'yelp': 711 | idx_train, idx_val, idx_test, c_num_mat = split_balance(torch.tensor(label_data), train_ratio=4, val_ratio=2, 712 | test_ratio=4) 713 | else: 714 | idx_train, idx_val, idx_test, c_num_mat = split_balance(torch.tensor(label_data), train_ratio=7, val_ratio=1, 715 | test_ratio=2) 716 | 717 | idx_train = torch.LongTensor(idx_train) 718 | idx_val = torch.LongTensor(idx_val) 719 | idx_test = torch.LongTensor(idx_test) 720 | 721 | labels = torch.LongTensor(np.where(label)[1]) 722 | adj = sparse_mx_to_torch_sparse_tensor(adj) 723 | 724 | return adj, features, labels, idx_train, idx_val, idx_test 725 | 726 | def sparse_mx_to_torch_sparse_tensor(sparse_mx): 727 | """Convert a scipy sparse matrix to a torch sparse tensor.""" 728 | sparse_mx = sparse_mx.tocoo().astype(np.float32) 729 | indices = torch.from_numpy( 730 | np.vstack((sparse_mx.row, sparse_mx.col)).astype(np.int64)) 731 | values = torch.from_numpy(sparse_mx.data) 732 | shape = torch.Size(sparse_mx.shape) 733 | return torch.sparse.FloatTensor(indices, values, shape) 734 | 735 | def encode_onehot(labels): 736 | 737 | classes = set(labels) 738 | classes_dict = {c: np.identity(len(classes))[i, :] for i, c in 739 | enumerate(classes)} 740 | labels_onehot = np.array(list(map(classes_dict.get, labels)), 741 | dtype=np.int32) 742 | 743 | return labels_onehot 744 | 745 | def normalize(mx): 746 | """Row-normalize sparse matrix""" 747 | rowsum = np.array(mx.sum(1)) 748 | r_inv = np.power(rowsum+1e-6, -1).flatten() 749 | r_inv[np.isinf(r_inv)] = 0. 750 | r_mat_inv = sp.diags(r_inv) 751 | mx = r_mat_inv.dot(mx) 752 | return mx 753 | -------------------------------------------------------------------------------- /utils/params.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # encoding: utf-8 3 | 4 | 5 | import argparse 6 | import torch 7 | import sys 8 | 9 | # argv = sys.argv 10 | # # dataset = argv[1] 11 | # # dataset = 'acm' 12 | # # dataset = 'aminer' 13 | 14 | 15 | def aminer_params(): 16 | parser = argparse.ArgumentParser() 17 | 18 | parser.add_argument('--node_entity_matching_path', type=str, default=r".\data\aminer_data\entityid_text_label.txt",help='path about the matching of node id and text id') 19 | parser.add_argument('--author_feature_path', type=str,default=r'.\data\aminer_data\id_author_feature.txt',help='path of aminer feature') 20 | parser.add_argument('--feature_path', type=str,default=r'.\data\aminer_data\keyword_feature.txt',help='path of aminer feature') 21 | parser.add_argument('--relation_path', type=str,default=r'.\data\aminer_data\relation.txt', help='path of aminer relationships') 22 | parser.add_argument('--pos_path', type=str, default=r'.\data\aminer_data\pos_index_5.txt',help='path of aminer pos label') 23 | parser.add_argument('--id_content_path', type=str, default=r'.\data\aminer_data\id_label_content.csv',help='to get the finetune features') 24 | 25 | parser.add_argument('--no-cuda', action='store_true', default=False, 26 | help='Disables CUDA training.') 27 | parser.add_argument('--fastmode', action='store_true', default=False, 28 | help='Validate during training pass.') 29 | parser.add_argument('--seed', type=int, default=50, help='Random seed.') 30 | parser.add_argument('--pretrain_epochs', type=int, default=200, 31 | help='Number of epochs to train.') 32 | parser.add_argument('--ft_epochs', type=int, default=200, 33 | help='Number of epochs to train during fine-tuning.') 34 | parser.add_argument('--lr', type=float, default=0.001, 35 | help='Initial learning rate.') 36 | parser.add_argument('--weight_decay', type=float, default=0.001, 37 | help='Weight decay (L2 loss on parameters).') 38 | parser.add_argument('--hidden', type=int, default=128, 39 | help='Number of hidden units.') 40 | parser.add_argument('--dropout', type=float, default=0.1, 41 | help='Dropout rate (1 - keep probability).') 42 | parser.add_argument('--node_dropout', type=float, default=0.1, 43 | help='Dropout rate (1 - keep probability).') 44 | parser.add_argument('--prune', default=True, help='network pruning for model pre-training') 45 | parser.add_argument('--prune_ratio', type=float, default=0.2, help='network pruning for model fine-tuning') 46 | 47 | parser.add_argument('--number_samples', type=int, default=1000, 48 | help='number of samples in co-modality pre-training') 49 | parser.add_argument('--device', type=str, default=torch.device("cuda:0" if torch.cuda.is_available() else "cpu"), 50 | help='use GPU') 51 | parser.add_argument('--finetune_device', type=str, default=torch.device("cuda:0" if torch.cuda.is_available() else "cpu"), 52 | help='use GPU') 53 | parser.add_argument('--image_encoder_model', type=str, default='swin_small_patch4_window7_224', 54 | help='resnet50, please refer to for more models') 55 | parser.add_argument('--text_encoder_model', type=str, default='distilbert-base-uncased', 56 | help='[bert-base-uncased,distilbert-base-uncased]') 57 | parser.add_argument('--text_encoder_tokenizer', type=str, default='distilbert-base-uncased', 58 | help='[bert-base-uncased,distilbert-base-uncased]') 59 | parser.add_argument('--max_length', type=int, default=70, help='the length of text') 60 | 61 | parser.add_argument('--node_encoder_model', type=str, default='gcn', help='[gcn,gat,sage,gin]') 62 | 63 | parser.add_argument('--nheads', type=int, default=8, help='gat') 64 | parser.add_argument('--alpha', type=float, default=0.2, help='gat') 65 | parser.add_argument('--image_size', type=int, default=224, help='the size of image') 66 | parser.add_argument('--imbalance_setting', type=str, default='focal', help='[reweight,upsample,focal]') 67 | parser.add_argument('--imbalance_up_scale', type=float, default=10.0, help='the scale for upsampling') 68 | 69 | parser.add_argument('--imbalance_ratio', type=float, default= 0, 70 | help='[0.01, 0.1, 0, 1] if 0 then original imbalance ratio if 1 then balanced data') 71 | parser.add_argument('--node_feature_dim', type=int, default=768, help='the dimension of aminer feature') 72 | parser.add_argument('--node_feature_project_dim', type=int, default=200, help='the dimension of projected feature') 73 | parser.add_argument('--node_embedding_dim', type=int, default=200, help='the dimension of node embedding') 74 | parser.add_argument('--hidden_dim', type=int, default=200) 75 | parser.add_argument('--out_dim', type=int, default=200) 76 | parser.add_argument('--nclass', type=int, default=5, help='the number of class') 77 | 78 | parser.add_argument('--num_projection_layers', type=int, default=1, 79 | help='the number of project layer for text/image and node') 80 | parser.add_argument('--feat_projection_dim', type=int, default=200, 81 | help='the dimension of projected feature for nodes') 82 | parser.add_argument('--projection_dim', type=int, default=256, 83 | help='the dimension of projected embedding of text/image and node') 84 | parser.add_argument('--pos', default=True) 85 | parser.add_argument('--debug', default=False) 86 | parser.add_argument('--batch_size', type=int, default=100,help='the size for each batch for co-modality pretraining') 87 | parser.add_argument('--num_workers', type=int, default=1, help='the number of workers') 88 | parser.add_argument('--patience', type=int, default=30, help='the number of epochs that the metric is not improved') 89 | parser.add_argument('--factor', type=int, default=0.5, help='the factor to change the learning rate') 90 | parser.add_argument('--temperature', type=int, default=0.1, help='the factor to change the learning rate') 91 | parser.add_argument('--pretrained', default=True, help="for text/image encoder") 92 | parser.add_argument('--trainable', default=False, help="for text/image encoder") 93 | parser.add_argument('--image_embedding_dim', type=int, default=768, help='the dimension of image embedding') 94 | parser.add_argument('--text_embedding_dim', type=int, default=768, help='the dimension of text embedding') 95 | parser.add_argument('--finetune_embedding_dim', type=int, default=768, help='the dimension of finetuen feature') 96 | parser.add_argument('--focal_alpha', type=int, default=0.75, help='focal loss') 97 | parser.add_argument('--focal_gamma', type=int, default=1.0, help='focal loss') 98 | parser.add_argument('--finetune', default=True, help='finetune the model') 99 | 100 | args, _ = parser.parse_known_args() 101 | 102 | return args 103 | 104 | def yelp_params(): 105 | parser = argparse.ArgumentParser() 106 | 107 | parser.add_argument('--node_entity_matching_path', type=str, default=r'',help='path about the matching of node id and text id') 108 | parser.add_argument('--feature_path', type=str, default=r'', help='path of yelp text feature') # base data here here with 0.25 109 | parser.add_argument('--embed_path', type=str, default=r'', 110 | help='path of yelp feature') 111 | parser.add_argument('--relation_path', type=str, default=r'',help='path of yelp relationships') 112 | parser.add_argument('--pos_path', type=str, default=r'', help='path of yelp pos label') 113 | parser.add_argument('--yelp_concate',default=False, help='whether concatenate yelp feature and embedding') 114 | parser.add_argument('--id_content_path', type=str, default=r'',help='to get the finetune features') 115 | 116 | parser.add_argument('--no-cuda', action='store_true', default=False, 117 | help='Disables CUDA training.') 118 | parser.add_argument('--fastmode', action='store_true', default=False, 119 | help='Validate during training pass.') 120 | parser.add_argument('--seed', type=int, default=50, help='Random seed.') 121 | parser.add_argument('--pretrain_epochs', type=int, default=30, 122 | help='Number of epochs to train.') 123 | parser.add_argument('--ft_epochs', type=int, default=60, 124 | help='Number of epochs to train.') 125 | parser.add_argument('--lr', type=float, default=0.01, 126 | help='Initial learning rate.') 127 | parser.add_argument('--weight_decay', type=float, default=0.001, 128 | help='Weight decay (L2 loss on parameters).') 129 | parser.add_argument('--hidden', type=int, default=128, 130 | help='Number of hidden units.') 131 | parser.add_argument('--dropout', type=float, default=0.5, 132 | help='Dropout rate (1 - keep probability).') 133 | parser.add_argument('--node_dropout', type=float, default=0.5, 134 | help='Dropout rate (1 - keep probability).') 135 | 136 | parser.add_argument('--prune', default=True, help='network pruning for model pre-training') 137 | parser.add_argument('--prune_ratio', type=float, default=0.5, help='network pruning for model fine-tuning') 138 | 139 | parser.add_argument('--number_samples', type=int, default=1000, 140 | help='number of samples in co-modality pre-training') 141 | parser.add_argument('--device', type=str, default=torch.device("cuda:0" if torch.cuda.is_available() else "cpu"), 142 | help='use GPU') 143 | parser.add_argument('--finetune_device', type=str, default=torch.device("cuda:0" if torch.cuda.is_available() else "cpu"), 144 | help='use GPU') 145 | parser.add_argument('--image_encoder_model', type=str, default='swin_small_patch4_window7_224', 146 | help='resnet50, please refer to for more models') 147 | parser.add_argument('--text_encoder_model', type=str, default='distilbert-base-uncased', 148 | help='[bert-base-uncased,distilbert-base-uncased]') 149 | parser.add_argument('--text_encoder_tokenizer', type=str, default='distilbert-base-uncased', 150 | help='[bert-base-uncased,distilbert-base-uncased]') 151 | parser.add_argument('--max_length', type=int, default=200, help='the length of text') 152 | parser.add_argument('--feat_projection_dim', type=int, default=200, 153 | help='the dimension of projected feature for nodes') 154 | parser.add_argument('--node_encoder_model', type=str, default='gcn', help='[gcn,gat,sage,gin]') 155 | parser.add_argument('--nheads', type=int, default=8, help='gat') 156 | parser.add_argument('--alpha', type=float, default=0.2, help='gat') 157 | parser.add_argument('--image_size', type=int, default=224, help='the size of image') 158 | parser.add_argument('--imbalance_setting', type=str, default='focal', help='[reweight,upsample,focal]') 159 | parser.add_argument('--imbalance_up_scale', type=float, default=10.0, help='the scale for upsampling') 160 | 161 | parser.add_argument('--imbalance_ratio', type=float, default=0, 162 | help='[0.01, 0.1, 0, 1] if 0 then original imbalance ratio if 1 then balanced data') 163 | 164 | parser.add_argument('--node_feature_dim', type=int, default=768, help='the dimension of yelp feature') 165 | parser.add_argument('--node_feature_project_dim', type=int, default=200, help='the dimension of projected feature') 166 | parser.add_argument('--node_embedding_dim', type=int, default=200, help='the dimension of node embedding') 167 | parser.add_argument('--hidden_dim', type=int, default=200) 168 | parser.add_argument('--out_dim', type=int, default=200) 169 | parser.add_argument('--nclass', type=int, default=2, help='the number of class') 170 | 171 | parser.add_argument('--num_projection_layers', type=int, default=1, 172 | help='the number of project layer for text/image and node') 173 | parser.add_argument('--projection_dim', type=int, default=256, 174 | help='the dimension of projected embedding of text/image and node') 175 | 176 | parser.add_argument('--pos', default=False) 177 | parser.add_argument('--debug', default=False) 178 | parser.add_argument('--batch_size', type=int, default=50, 179 | help='the size for each batch for co-modality pretraining') 180 | parser.add_argument('--num_workers', type=int, default=1, help='the number of workers') 181 | parser.add_argument('--patience', type=int, default=1, help='the number of epochs that the metric is not improved') 182 | parser.add_argument('--factor', type=int, default=0.5, help='the factor to change the learning rate') 183 | parser.add_argument('--temperature', type=int, default=1.0, help='the factor to change the learning rate') 184 | 185 | parser.add_argument('--pretrained', default=True,help = "for text/image encoder") 186 | parser.add_argument('--trainable', default=False, help = "for text/image encoder") 187 | 188 | parser.add_argument('--image_embedding_dim', type=int, default=768, help='the dimension of image embedding') 189 | parser.add_argument('--text_embedding_dim', type=int, default=768, help='the dimension of text embedding') 190 | parser.add_argument('--finetune_embedding_dim', type=int, default=768, help='the dimension of finetuen feature') 191 | parser.add_argument('--focal_alpha', type=int, default=0.5, help='focal loss') 192 | parser.add_argument('--focal_gamma', type=int, default=0.6, help='focal loss') 193 | parser.add_argument('--finetune', default=True, help='finetune the model') 194 | 195 | 196 | args, _ = parser.parse_known_args() 197 | 198 | return args 199 | 200 | def github_params(): 201 | parser = argparse.ArgumentParser() 202 | 203 | parser.add_argument('--node_entity_matching_path', type=str,default=r'',help='path about the matching of node id and text id') 204 | parser.add_argument('--feature_path', type=str, default=r'', 205 | help='path of github feature') 206 | parser.add_argument('--relation_path', type=str, default=r'', 207 | help='path of github relationships') 208 | parser.add_argument('--pos_path', type=str, default=r'', 209 | help='path of github pos label') 210 | parser.add_argument('--label_path', type=str, default=r'', 211 | help='path of label') 212 | 213 | parser.add_argument('--no-cuda', action='store_true', default=False, 214 | help='Disables CUDA training.') 215 | parser.add_argument('--fastmode', action='store_true', default=False, 216 | help='Validate during training pass.') 217 | parser.add_argument('--seed', type=int, default=50, help='Random seed.') 218 | parser.add_argument('--pretrain_epochs', type=int, default=60, 219 | help='Number of epochs to train.') 220 | parser.add_argument('--ft_epochs', type=int, default=60, 221 | help='Number of epochs to train.') 222 | parser.add_argument('--lr', type=float, default=0.01, 223 | help='Initial learning rate.') 224 | parser.add_argument('--weight_decay', type=float, default=0.001, 225 | help='Weight decay (L2 loss on parameters).') 226 | parser.add_argument('--hidden', type=int, default=128, 227 | help='Number of hidden units.') 228 | parser.add_argument('--dropout', type=float, default=0.5, 229 | help='Dropout rate (1 - keep probability).') 230 | parser.add_argument('--node_dropout', type=float, default=0.5, 231 | help='Dropout rate (1 - keep probability).') 232 | 233 | parser.add_argument('--prune', default=True, help='network pruning for model pre-training') 234 | parser.add_argument('--prune_ratio', type=float, default=0.5, help='network pruning for model fine-tuning') 235 | 236 | parser.add_argument('--number_samples', type=int, default=1000, 237 | help='number of samples in co-modality pre-training') 238 | parser.add_argument('--device', type=str, default=torch.device("cuda:0" if torch.cuda.is_available() else "cpu"), 239 | help='use GPU') 240 | parser.add_argument('--finetune_device', type=str, default=torch.device("cuda:0" if torch.cuda.is_available() else "cpu"), 241 | help='use GPU') 242 | parser.add_argument('--id_content_path', type=str, 243 | default=r'', 244 | help='to get the finetune features') 245 | parser.add_argument('--image_encoder_model', type=str, default='swin_small_patch4_window7_224', 246 | help='resnet50, please refer to for more models') 247 | parser.add_argument('--text_encoder_model', type=str, default='distilbert-base-uncased', 248 | help='[bert-base-uncased,distilbert-base-uncased]') 249 | parser.add_argument('--text_encoder_tokenizer', type=str, default='distilbert-base-uncased', 250 | help='[bert-base-uncased,distilbert-base-uncased]') 251 | parser.add_argument('--max_length', type=int, default=200, help='the length of text') 252 | parser.add_argument('--feat_projection_dim', type=int, default=200, 253 | help='the dimension of projected feature for nodes') 254 | parser.add_argument('--node_encoder_model', type=str, default='gcn', help='[gcn,gat,sage,gin]') 255 | parser.add_argument('--nheads', type=int, default=8, help='gat') 256 | parser.add_argument('--alpha', type=float, default=0.2, help='gat') 257 | parser.add_argument('--image_size', type=int, default=224, help='the size of image') 258 | parser.add_argument('--imbalance_setting', type=str, default='focal', help='[reweight,upsample,focal]') 259 | parser.add_argument('--imbalance_up_scale', type=float, default=10.0, help='the scale for upsampling') 260 | 261 | parser.add_argument('--imbalance_ratio', type=float, default=0, 262 | help='[0.01, 0.1, 0, 1] if 0 then original imbalance ratio if 1 then balanced data') 263 | 264 | parser.add_argument('--node_feature_dim', type=int, default=128, help='the dimension of github feature') 265 | parser.add_argument('--node_feature_project_dim', type=int, default=200, help='the dimension of projected feature') 266 | parser.add_argument('--node_embedding_dim', type=int, default=200, help='the dimension of node embedding') 267 | parser.add_argument('--hidden_dim', type=int, default=200) 268 | parser.add_argument('--out_dim', type=int, default=200) 269 | parser.add_argument('--nclass', type=int, default=2, help='the number of class') 270 | 271 | parser.add_argument('--num_projection_layers', type=int, default=1, 272 | help='the number of project layer for text/image and node') 273 | parser.add_argument('--projection_dim', type=int, default=256, 274 | help='the dimension of projected embedding of text/image and node') 275 | 276 | parser.add_argument('--pos', default=False) 277 | parser.add_argument('--debug', default=False) 278 | parser.add_argument('--batch_size', type=int, default=80, 279 | help='the size for each batch for co-modality pretraining') 280 | parser.add_argument('--num_workers', type=int, default=1, help='the number of workers') 281 | parser.add_argument('--patience', type=int, default=5, help='the number of epochs that the metric is not improved') 282 | parser.add_argument('--factor', type=int, default=0.5, help='the factor to change the learning rate') 283 | parser.add_argument('--temperature', type=int, default=1.0, help='the factor to change the learning rate') 284 | 285 | parser.add_argument('--pretrained', default=True,help = "for text/image encoder") 286 | parser.add_argument('--trainable', default=True, help = "for text/image encoder") 287 | parser.add_argument('--image_embedding_dim', type=int, default=768, help='the dimension of image embedding') 288 | parser.add_argument('--text_embedding_dim', type=int, default=768, help='the dimension of text embedding') 289 | parser.add_argument('--finetune_embedding_dim', type=int, default=768, help='the dimension of finetuen feature') 290 | parser.add_argument('--focal_alpha', type=int, default=0.25, help='focal loss') 291 | parser.add_argument('--focal_gamma', type=int, default=2.0, help='focal loss') 292 | parser.add_argument('--finetune', default=True, help='finetune the model') 293 | 294 | args, _ = parser.parse_known_args() 295 | 296 | return args 297 | 298 | def instagram_params(): 299 | parser = argparse.ArgumentParser() 300 | 301 | parser.add_argument('--node_entity_matching_path', type=str, 302 | default=r'', 303 | help='path about the matching of node id and picture id') 304 | parser.add_argument('--feature_path', type=str, default=r'', 305 | help='path of instagram feature') 306 | parser.add_argument('--relation_path', type=str, 307 | default=r'', 308 | help='path of instagram relationships') 309 | parser.add_argument('--pos_path', type=str, default=r'', 310 | help='path of instagram pos label') 311 | parser.add_argument('--id_content_path', type=str, default=r'', 312 | help='to get the finetune features') 313 | 314 | parser.add_argument('--no-cuda', action='store_true', default=False, 315 | help='Disables CUDA training.') 316 | parser.add_argument('--fastmode', action='store_true', default=False, 317 | help='Validate during training pass.') 318 | parser.add_argument('--seed', type=int, default=50, help='Random seed.') 319 | parser.add_argument('--pretrain_epochs', type=int, default=1, 320 | help='Number of epochs to train.') # 60 321 | parser.add_argument('--ft_epochs', type=int, default=60, 322 | help='Number of epochs to train.') 323 | parser.add_argument('--lr', type=float, default=0.01, 324 | help='Initial learning rate.') 325 | parser.add_argument('--weight_decay', type=float, default=0.001, 326 | help='Weight decay (L2 loss on parameters).') 327 | parser.add_argument('--hidden', type=int, default=128, 328 | help='Number of hidden units.') 329 | parser.add_argument('--dropout', type=float, default=0.5, 330 | help='Dropout rate (1 - keep probability).') 331 | parser.add_argument('--node_dropout', type=float, default=0.5, 332 | help='Dropout rate (1 - keep probability).') 333 | 334 | parser.add_argument('--prune', default=True, help='network pruning for model pre-training') 335 | parser.add_argument('--prune_ratio', type=float, default=0.5, help='network pruning for model fine-tuning') 336 | 337 | parser.add_argument('--number_samples', type=int, default=1000, 338 | help='number of samples in co-modality pre-training') 339 | parser.add_argument('--device', type=str, default=torch.device("cuda:0" if torch.cuda.is_available() else "cpu"), 340 | help='use GPU') 341 | parser.add_argument('--finetune_device', type=str, default=torch.device("cuda:0" if torch.cuda.is_available() else "cpu"), 342 | help='use GPU') 343 | parser.add_argument('--image_encoder_model', type=str, default='swin_small_patch4_window7_224', 344 | help='resnet50, please refer to for more models') 345 | parser.add_argument('--text_encoder_model', type=str, default='distilbert-base-uncased', 346 | help='[bert-base-uncased,distilbert-base-uncased]') 347 | parser.add_argument('--text_encoder_tokenizer', type=str, default='distilbert-base-uncased', 348 | help='[bert-base-uncased,distilbert-base-uncased]') 349 | parser.add_argument('--max_length', type=int, default=200, help='the length of text') 350 | parser.add_argument('--feat_projection_dim', type=int, default=200, 351 | help='the dimension of projected feature for nodes') 352 | 353 | parser.add_argument('--node_encoder_model', type=str, default='gcn', help='[gcn,gat,sage,gin]') 354 | parser.add_argument('--nheads', type=int, default=8, help='gat') 355 | parser.add_argument('--alpha', type=float, default=0.2, help='gat') 356 | 357 | parser.add_argument('--imbalance_setting', type=str, default='focal', help='[reweight,upsample,focal]') 358 | parser.add_argument('--imbalance_up_scale', type=float, default=10.0, help='the scale for upsampling') 359 | 360 | parser.add_argument('--imbalance_ratio', type=float, default=0, 361 | help='[0.01, 0.1, 0, 1] if 0 then original imbalance ratio if 1 then balanced data') 362 | 363 | parser.add_argument('--node_feature_dim', type=int, default=768, help='the dimension of instagram feature') 364 | parser.add_argument('--node_feature_project_dim', type=int, default=200, help='the dimension of projected feature') 365 | parser.add_argument('--node_embedding_dim', type=int, default=200, help='the dimension of node embedding') 366 | parser.add_argument('--hidden_dim', type=int, default=200) 367 | parser.add_argument('--out_dim', type=int, default=200) 368 | parser.add_argument('--nclass', type=int, default=2, help='the number of class') 369 | 370 | parser.add_argument('--image_size', type=int, default=224, help='the size of image') 371 | parser.add_argument('--image_embedding_dim', type=int, default=768, help='the dimension of image embedding') 372 | parser.add_argument('--text_embedding_dim', type=int, default=768, help='the dimension of text embedding') 373 | parser.add_argument('--finetune_embedding_dim', type=int, default=768, help='the dimension of finetuen feature') 374 | 375 | parser.add_argument('--num_projection_layers', type=int, default=1, 376 | help='the number of project layer for text/image and node') 377 | parser.add_argument('--projection_dim', type=int, default=256, 378 | help='the dimension of projected embedding of text/image and node') 379 | 380 | parser.add_argument('--pos', default=False) 381 | parser.add_argument('--debug', default=False) 382 | parser.add_argument('--batch_size', type=int, default=30, 383 | help='the size for each batch for co-modality pretraining') 384 | parser.add_argument('--num_workers', type=int, default=2, help='the number of workers') 385 | parser.add_argument('--patience', type=int, default=5, help='the number of epochs that the metric is not improved') 386 | parser.add_argument('--factor', type=int, default=0.5, help='the factor to change the learning rate') 387 | parser.add_argument('--temperature', type=int, default=1.0, help='the factor to change the learning rate') 388 | 389 | parser.add_argument('--pretrained', default=True,help = "for text/image encoder") 390 | parser.add_argument('--trainable', default=True, help = "for text/image encoder") 391 | 392 | parser.add_argument('--focal_alpha', type=int, default=0.25, help='focal loss') 393 | parser.add_argument('--focal_gamma', type=int, default=1.0, help='focal loss') 394 | parser.add_argument('--finetune', default=True, help='finetune the model') 395 | 396 | args, _ = parser.parse_known_args() 397 | 398 | return args 399 | 400 | 401 | def set_params(dataset): 402 | if dataset == "aminer": 403 | args = aminer_params() 404 | elif dataset == "yelp": 405 | args = yelp_params() 406 | elif dataset == "github": 407 | args = github_params() 408 | elif dataset == "instagram": 409 | args = instagram_params() 410 | 411 | args.dataset = dataset 412 | 413 | return args 414 | 415 | 416 | dataset = 'aminer' 417 | # dataset = 'yelp' 418 | # dataset = 'github' 419 | # dataset = 'instagram' 420 | args = set_params(dataset) 421 | -------------------------------------------------------------------------------- /utils/util.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import numpy as np 4 | import torch 5 | import torch.nn.functional as F 6 | from sklearn.metrics import roc_auc_score, f1_score 7 | 8 | 9 | class AvgMeter: 10 | def __init__(self, name="Metric"): 11 | self.name = name 12 | self.reset() 13 | 14 | def reset(self): 15 | self.avg, self.sum, self.count = [0] * 3 16 | 17 | def update(self, val, count=1): 18 | self.count += count 19 | self.sum += val * count 20 | self.avg = self.sum / self.count 21 | 22 | def __repr__(self): 23 | text = f"{self.name}: {self.avg:.4f}" 24 | return text 25 | 26 | 27 | def get_lr(optimizer): 28 | for param_group in optimizer.param_groups: 29 | return param_group["lr"] 30 | 31 | 32 | def seed_torch(seed=1029): 33 | random.seed(seed) 34 | os.environ['PYTHONHASHSEED'] = str(seed) 35 | np.random.seed(seed) 36 | torch.manual_seed(seed) 37 | torch.cuda.manual_seed(seed) 38 | torch.cuda.manual_seed_all(seed) # if you are using multi-GPU. 39 | torch.backends.cudnn.benchmark = False 40 | torch.backends.cudnn.deterministic = True 41 | 42 | 43 | 44 | def evaluate_performance(labels, output): 45 | f1_micro = f1_score(labels.cpu().detach().numpy(), 46 | output.max(1)[1].cpu().detach().numpy(), average='micro') 47 | f1_macro = f1_score(labels.cpu().detach().numpy(), 48 | output.max(1)[1].cpu().detach().numpy(), average='macro') 49 | 50 | if labels.max() > 1: 51 | auc = roc_auc_score(labels.detach().cpu(), 52 | F.softmax(output, dim=-1).detach().cpu(), average='macro', 53 | multi_class='ovr') 54 | else: 55 | 56 | auc = roc_auc_score(labels.detach().cpu(), 57 | F.softmax(output, dim=-1)[:, 1].detach().cpu(), average='macro') 58 | 59 | # auc = roc_auc_score(labels.detach().cpu(), 60 | # torch.nan_to_num(F.softmax(output, dim=-1)[:, 1], 1e-5).detach().cpu(), average='macro') 61 | 62 | return f1_micro, f1_macro, auc 63 | --------------------------------------------------------------------------------