-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathn_pretrain_contextpred.py
174 lines (131 loc) · 7.79 KB
/
n_pretrain_contextpred.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
import argparse
from loader import BioDataset
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from tqdm import tqdm
import numpy as np
from model_node import GNN
from sklearn.metrics import roc_auc_score
import pandas as pd
from util import ExtractSubstructureContextPair
from dataloader import DataLoaderSubstructContext
from torch_geometric.nn import global_add_pool, global_mean_pool, global_max_pool
def pool_func(x, batch, mode = "sum"):
if mode == "sum":
return global_add_pool(x, batch)
elif mode == "mean":
return global_mean_pool(x, batch)
elif mode == "max":
return global_max_pool(x, batch)
def cycle_index(num, shift):
arr = torch.arange(num) + shift
arr[-shift:] = torch.arange(shift)
return arr
criterion = nn.BCEWithLogitsLoss()
def train(args, model_substruct, model_context, loader, optimizer_substruct, optimizer_context, device):
model_substruct.train()
balanced_loss_accum = 0
acc_accum = 0
for step, batch in enumerate(tqdm(loader, desc="Iteration")):
batch = batch.to(device)
#print(batch)
# creating substructure representation
substruct_rep = model_substruct(batch.x_substruct, batch.edge_index_substruct, batch.edge_attr_substruct)[batch.center_substruct_idx]
### creating context representations
overlapped_node_rep = model_context(batch.x_context, batch.edge_index_context, batch.edge_attr_context)[batch.overlap_context_substruct_idx]
#Contexts are represented by
if args.mode == "cbow":
# positive context representation
context_rep = pool_func(overlapped_node_rep, batch.batch_overlapped_context, mode = args.context_pooling)
# negative contexts are obtained by shifting the indicies of context embeddings
neg_context_rep = torch.cat([context_rep[cycle_index(len(context_rep), i+1)] for i in range(args.neg_samples)], dim = 0)
pred_pos = torch.sum(substruct_rep * context_rep, dim = 1)
pred_neg = torch.sum(substruct_rep.repeat((args.neg_samples, 1))*neg_context_rep, dim = 1)
elif args.mode == "skipgram":
expanded_substruct_rep = torch.cat([substruct_rep[i].repeat((batch.overlapped_context_size[i],1)) for i in range(len(substruct_rep))], dim = 0)
pred_pos = torch.sum(expanded_substruct_rep * overlapped_node_rep, dim = 1)
#shift indices of substructures to create negative examples
shifted_expanded_substruct_rep = []
for i in range(args.neg_samples):
shifted_substruct_rep = substruct_rep[cycle_index(len(substruct_rep), i+1)]
shifted_expanded_substruct_rep.append(torch.cat([shifted_substruct_rep[i].repeat((batch.overlapped_context_size[i],1)) for i in range(len(shifted_substruct_rep))], dim = 0))
shifted_expanded_substruct_rep = torch.cat(shifted_expanded_substruct_rep, dim = 0)
pred_neg = torch.sum(shifted_expanded_substruct_rep * overlapped_node_rep.repeat((args.neg_samples, 1)), dim = 1)
else:
raise ValueError("Invalid mode!")
loss_pos = criterion(pred_pos.double(), torch.ones(len(pred_pos)).to(pred_pos.device).double())
loss_neg = criterion(pred_neg.double(), torch.zeros(len(pred_neg)).to(pred_neg.device).double())
optimizer_substruct.zero_grad()
optimizer_context.zero_grad()
loss = loss_pos + args.neg_samples*loss_neg
loss.backward()
#To write: optimizer
optimizer_substruct.step()
optimizer_context.step()
balanced_loss_accum += float(loss_pos.detach().cpu().item() + loss_neg.detach().cpu().item())
acc_accum += 0.5* (float(torch.sum(pred_pos > 0).detach().cpu().item())/len(pred_pos) + float(torch.sum(pred_neg < 0).detach().cpu().item())/len(pred_neg))
return balanced_loss_accum/step, acc_accum/step
def main():
# Training settings
parser = argparse.ArgumentParser(description='PyTorch implementation of pre-training of graph neural networks')
parser.add_argument('--device', type=int, default=0,
help='which gpu to use if any (default: 0)')
parser.add_argument('--batch_size', type=int, default=256,
help='input batch size for training (default: 256)')
parser.add_argument('--epochs', type=int, default=100,
help='number of epochs to train (default: 100)')
parser.add_argument('--lr', type=float, default=0.001,
help='learning rate (default: 0.001)')
parser.add_argument('--decay', type=float, default=0,
help='weight decay (default: 0)')
parser.add_argument('--num_layer', type=int, default=5,
help='number of GNN message passing layers (default: 5).')
parser.add_argument('--l1', type=int, default=1,
help='l1 (default: 1).')
parser.add_argument('--center', type=int, default=0,
help='center (default: 0).')
parser.add_argument('--emb_dim', type=int, default=300,
help='embedding dimensions (default: 300)')
parser.add_argument('--dropout_ratio', type=float, default=0,
help='dropout ratio (default: 0)')
parser.add_argument('--neg_samples', type=int, default=1,
help='number of negative contexts per positive context (default: 1)')
parser.add_argument('--JK', type=str, default="last",
help='how the node features are combined across layers. last, sum, max or concat')
parser.add_argument('--context_pooling', type=str, default="mean",
help='how the contexts are pooled (sum, mean, or max)')
parser.add_argument('--gnn_type', type=str, default="gin")
parser.add_argument('--mode', type=str, default = "cbow", help = "cbow or skipgram")
parser.add_argument('--model_file', type=str, default = '', help='filename to output the model')
parser.add_argument('--num_workers', type=int, default = 4, help='number of workers for dataset loading')
args = parser.parse_args()
torch.manual_seed(0)
np.random.seed(0)
device = torch.device("cuda:" + str(args.device)) if torch.cuda.is_available() else torch.device("cpu")
if torch.cuda.is_available():
torch.cuda.manual_seed_all(0)
print(args.mode)
#set up dataset
root_unsupervised = 'dataset/unsupervised'
dataset = BioDataset(root_unsupervised, data_type='unsupervised', transform = ExtractSubstructureContextPair(l1 = args.l1, center = args.center))
print(dataset[0])
print("l1: " + str(args.l1))
print("center: " + str(args.center))
loader = DataLoaderSubstructContext(dataset, batch_size=args.batch_size, shuffle=True, num_workers = args.num_workers)
#print(dataset[0])
#set up models, one for pre-training and one for context embeddings
model_substruct = GNN(args.num_layer, args.emb_dim, JK = args.JK, drop_ratio = args.dropout_ratio, gnn_type = args.gnn_type).to(device)
model_context = GNN(3, args.emb_dim, JK = args.JK, drop_ratio = args.dropout_ratio, gnn_type = args.gnn_type).to(device)
#set up optimizer for the two GNNs
optimizer_substruct = optim.Adam(model_substruct.parameters(), lr=args.lr, weight_decay=args.decay)
optimizer_context = optim.Adam(model_context.parameters(), lr=args.lr, weight_decay=args.decay)
for epoch in range(1, args.epochs+1):
print("====epoch " + str(epoch))
train_loss, train_acc = train(args, model_substruct, model_context, loader, optimizer_substruct, optimizer_context, device)
print(train_loss, train_acc)
if not args.model_file == "":
torch.save(model_substruct.state_dict(), args.model_file + ".pth")
if __name__ == "__main__":
main()