89 lines
3.3 KiB
Python
89 lines
3.3 KiB
Python
from collections import defaultdict
|
||
|
||
import numpy as np
|
||
import torch
|
||
|
||
|
||
def remove_unseen_classes_from_training(train_mask, labels, removed_class):
|
||
"""Remove the unseen classes (the first three classes by default) to get the zero-shot (i.e., completely imbalanced) label setting
|
||
Input: train_mask, labels, removed_class
|
||
Output: train_mask_zs: the bool list only containing seen classes
|
||
"""
|
||
train_mask_zs = train_mask.clone()
|
||
for i in range(train_mask_zs.numel()):
|
||
if train_mask_zs[i] == 1 and (labels[i].item() in removed_class):
|
||
train_mask_zs[i] = 0
|
||
return train_mask_zs
|
||
|
||
|
||
def get_class_set(labels):
|
||
"""Get the class set.
|
||
Input: labels [l, [c1, c2, ..]]
|
||
Output:the labeled class set dict_keys([k1, k2, ..])
|
||
"""
|
||
mydict = {}
|
||
for y in labels:
|
||
for label in y:
|
||
mydict[int(label)] = 1
|
||
return mydict.keys()
|
||
|
||
|
||
def get_label_attributes(train_mask_zs, nodeids, labellist, features):
|
||
"""Get the class-center (semanic knowledge) of each seen class.
|
||
Suppose a node i is labeled as c, then attribute[c] += node_i_attribute, finally mean(attribute[c])
|
||
Input: train_mask_zs, nodeids, labellist, features
|
||
Output: label_attribute{}: label -> average_labeled_node_features (class centers)
|
||
"""
|
||
_, feat_num = features.shape
|
||
labels = get_class_set(labellist)
|
||
label_attribute_nodes = defaultdict(list)
|
||
for nodeid, labels in zip(nodeids, labellist):
|
||
for label in labels:
|
||
label_attribute_nodes[int(label)].append(int(nodeid))
|
||
label_attribute = {}
|
||
for label in label_attribute_nodes.keys():
|
||
nodes = label_attribute_nodes[int(label)]
|
||
selected_features = features[nodes, :]
|
||
label_attribute[int(label)] = np.mean(selected_features, axis=0)
|
||
return label_attribute
|
||
|
||
|
||
def get_labeled_nodes_label_attribute(train_mask_zs, labels, features, cuda):
|
||
"""Replace the original labels by their class-centers.
|
||
For each label c in the training set, the following operations will be performed:
|
||
Get label_attribute{} through function get_label_attributes, then res[i, :] = label_attribute[c]
|
||
Input: train_mask_zs, labels, features
|
||
Output: Y_{semantic} [l, ft]: tensor
|
||
"""
|
||
X = torch.LongTensor(range(features.shape[0]))
|
||
nodeids = []
|
||
labellist = []
|
||
for i in X[train_mask_zs].numpy().tolist():
|
||
nodeids.append(str(i))
|
||
for i in labels[train_mask_zs].cpu().numpy().tolist():
|
||
labellist.append([str(i)])
|
||
|
||
# 1. get the semantic knowledge (class centers) of all seen classes
|
||
label_attribute = get_label_attributes(
|
||
train_mask_zs=train_mask_zs,
|
||
nodeids=nodeids,
|
||
labellist=labellist,
|
||
features=features.cpu().numpy(),
|
||
)
|
||
|
||
# 2. replace original labels by their class centers (semantic knowledge)
|
||
res = np.zeros([len(nodeids), features.shape[1]])
|
||
for i, labels in enumerate(labellist):
|
||
# support mutiple labels
|
||
c = len(labels)
|
||
temp = np.zeros([c, features.shape[1]])
|
||
for ii, label in enumerate(labels):
|
||
temp[ii, :] = label_attribute[int(label)]
|
||
temp = np.mean(temp, axis=0)
|
||
res[i, :] = temp
|
||
if cuda:
|
||
res = torch.FloatTensor(res).cuda()
|
||
else:
|
||
res = torch.FloatTensor(res)
|
||
return res
|