class DAEGC(nn.Module):
def __init__(self,
num_features,
hidden_size,
embedding_size,
alpha,
num_clusters,
v=1):
super(DAEGC, self).__init__()
self.num_clusters = num_clusters
self.v = v
# get pretrain model
self.gat = GAT(num_features, hidden_size, embedding_size, alpha)
self.gat.load_state_dict(
torch.load(args.pretrain_path, map_location='cpu'))
# cluster layer
self.cluster_layer = Parameter(
torch.Tensor(num_clusters, embedding_size))
torch.nn.init.xavier_normal_(self.cluster_layer.data)
def forward(self, x, adj, M):
A_pred, z = self.gat(x, adj, M)
q = self.get_Q(z)
return A_pred, z, q
def get_Q(self, z):
q = 1.0 / (1.0 + torch.sum(
torch.pow(z.unsqueeze(1) - self.cluster_layer, 2), 2) / self.v)
q = q.pow((self.v + 1.0) / 2.0)
q = (q.t() / torch.sum(q, 1)).t()
return q
if loss_values[-1] < best:
best = loss_values[-1]
best_epoch = epoch
patience_counter = 0
else:
patience_counter += 1
if patience_counter == args.patience:
break
files = glob.glob('*.pkl')
for file in files:
epoch_nb = int(file.split('.')[0])
if epoch_nb < best_epoch:
os.remove(file)
files = glob.glob('*.pkl')
for file in files:
epoch_nb = int(file.split('.')[0])
if epoch_nb > best_epoch:
os.remove(file)
print("Optimization Finished!")
print("Total time elapsed: {:.4f}s".format(time.time() - start_time))
# Restore best model
print('Loading {}th epoch'.format(best_epoch))
model.load_state_dict(torch.load('{}.pkl'.format(best_epoch)))
# Testing
run_test(model, features, adj, idx_test, labels)
else:
bad_counter += 1
if bad_counter == args.patience:
break
files = glob.glob('*' + save_name)
for file in files:
epoch_nb = int(file.split('.')[0])
if epoch_nb < best_epoch:
os.remove(file)
files = glob.glob('*' + save_name)
for file in files:
epoch_nb = int(file.split('.')[0])
if epoch_nb > best_epoch:
os.remove(file)
print("Optimization Finished!")
print("Total time elapsed: {:.4f}s".format(time.time() - t_total))
# Restore best model
print('Loading {}th epoch'.format(best_epoch))
model.load_state_dict(torch.load(('{}' + save_name).format(best_epoch)))
# Testing
compute_test()
# print(loss_values)
print('OK.')
def main(args):
# load and preprocess dataset
data = CoraGraphDataset()
g = data[0]
if args.gpu < 0:
cuda = False
else:
cuda = True
g = g.int().to(args.gpu)
features = g.ndata['feat']
labels = g.ndata['label']
train_mask = g.ndata['train_mask']
val_mask = g.ndata['val_mask']
test_mask = g.ndata['test_mask']
num_feats = features.shape[1]
n_classes = data.num_labels
n_edges = data.graph.number_of_edges()
x = g.nodes().cpu()
print(features)
print(g)
x = input()
print("""----Data statistics------'
#Edges %d
#Classes %d
#Train samples %d
#Val samples %d
#Test samples %d""" %
(n_edges, n_classes,
train_mask.int().sum().item(),
val_mask.int().sum().item(),
test_mask.int().sum().item()))
# add self loop
g = dgl.remove_self_loop(g)
g = dgl.add_self_loop(g)
n_edges = g.number_of_edges()
# create model
heads = ([args.num_heads] * args.num_layers) + [args.num_out_heads]
model = GAT(args.num_layers,
num_feats,
args.num_hidden,
n_classes,
heads,
F.elu,
args.in_drop,
args.attn_drop,
args.negative_slope,
args.residual)
print(model)
if args.early_stop:
stopper = EarlyStopping(patience=100)
if cuda:
model.cuda()
loss_fcn = torch.nn.CrossEntropyLoss()
# use optimizer
optimizer = torch.optim.Adam(
model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
# initialize graph
dur = []
for epoch in range(args.epochs):
model.train()
if epoch >= 3:
t0 = time.time()
# forward
logits = model(g, features)
loss = loss_fcn(logits[train_mask], labels[train_mask])
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch >= 3:
dur.append(time.time() - t0)
train_acc = accuracy(logits[train_mask], labels[train_mask])
if args.fastmode:
val_acc = accuracy(logits[val_mask], labels[val_mask])
else:
val_acc = evaluate(model, g, features, labels, val_mask)
if args.early_stop:
if stopper.step(val_acc, model):
break
print("Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | TrainAcc {:.4f} |"
" ValAcc {:.4f} | ETputs(KTEPS) {:.2f}".
format(epoch, np.mean(dur), loss.item(), train_acc,
val_acc, n_edges / np.mean(dur) / 1000))
print()
if args.early_stop:
model.load_state_dict(torch.load('es_checkpoint.pt'))
acc = evaluate(model, g, features, labels, test_mask)
print("Test Accuracy {:.4f}".format(acc))
def test(checkpoint_path, class_num, args):
model = GAT(args.feat_dim, args.embed_dim, class_num, args.alpha,
args.dropout, args.nheads, args.use_cuda)
checkpoint = torch.load(checkpoint_path)
model.load_state_dict(checkpoint['graph_state_dict'])
if args.use_cuda:
model.cuda()
model.eval()
for key in building.keys():
node_num = test_dataset[key]['node_num']
old_feature_map, adj_lists = collectGraph_test(
test_dataset[key]['feature_path'], node_num, args.feat_dim,
args.num_sample, args.suffix)
old_feature_map = torch.FloatTensor(old_feature_map)
if args.use_cuda:
old_feature_map = old_feature_map.cuda()
batch_num = int(math.ceil(node_num / float(args.batch_size)))
new_feature_map = torch.FloatTensor()
for batch in tqdm(range(batch_num)):
start_node = batch * args.batch_size
end_node = min((batch + 1) * args.batch_size, node_num)
batch_nodes = range(start_node, end_node)
batch_neighbors = [adj_lists[node] for node in batch_nodes]
new_feature, _ = model(old_feature_map, batch_nodes,
batch_neighbors)
new_feature = F.normalize(new_feature, p=2, dim=1)
new_feature_map = torch.cat(
(new_feature_map, new_feature.cpu().detach()), dim=0)
new_feature_map = new_feature_map.numpy()
old_similarity = np.dot(old_feature_map.cpu().numpy(),
old_feature_map.cpu().numpy().T)
new_similarity = np.dot(new_feature_map, new_feature_map.T)
mAP_old = building[key].evalRetrieval(old_similarity, retrieval_result)
mAP_new = building[key].evalRetrieval(new_similarity, retrieval_result)
print time.strftime('%Y-%m-%d %H:%M:%S'), 'eval {}'.format(key)
print 'base feature: {}, new feature: {}'.format(
old_feature_map.size(), new_feature_map.shape)
print 'base mAP: {:.4f}, new mAP: {:.4f}, improve: {:.4f}'.format(
mAP_old, mAP_new, mAP_new - mAP_old)
## directly update node's features by mean pooling features of its neighbors.
meanAggregator = model.attentions[0]
mean_feature_map = torch.FloatTensor()
for batch in tqdm(range(batch_num)):
start_node = batch * args.batch_size
end_node = min((batch + 1) * args.batch_size, node_num)
batch_nodes = range(start_node, end_node)
batch_neighbors = [adj_lists[node] for node in batch_nodes]
mean_feature = meanAggregator.meanAggregate(
old_feature_map, batch_nodes, batch_neighbors)
mean_feature = F.normalize(mean_feature, p=2, dim=1)
mean_feature_map = torch.cat(
(mean_feature_map, mean_feature.cpu().detach()), dim=0)
mean_feature_map = mean_feature_map.numpy()
mean_similarity = np.dot(mean_feature_map, mean_feature_map.T)
mAP_mean = building[key].evalRetrieval(mean_similarity,
retrieval_result)
print 'mean aggregation mAP: {:.4f}'.format(mAP_mean)
print ""
