wehub-resource-sync
34 lines
1.1 KiB
Python
34 lines
1.1 KiB
Python
import dgl
|
|
import mxnet as mx
|
|
import numpy as np
|
|
from dgl.utils import toindex
|
|
|
|
|
|
def l0_sample(g, positive_max=128, negative_ratio=3):
|
|
"""sampling positive and negative edges"""
|
|
if g is None:
|
|
return None
|
|
n_eids = g.number_of_edges()
|
|
pos_eids = np.where(g.edata["rel_class"].asnumpy() > 0)[0]
|
|
neg_eids = np.where(g.edata["rel_class"].asnumpy() == 0)[0]
|
|
if len(pos_eids) == 0:
|
|
return None
|
|
|
|
positive_num = min(len(pos_eids), positive_max)
|
|
negative_num = min(len(neg_eids), positive_num * negative_ratio)
|
|
pos_sample = np.random.choice(pos_eids, positive_num, replace=False)
|
|
neg_sample = np.random.choice(neg_eids, negative_num, replace=False)
|
|
weights = np.zeros(n_eids)
|
|
# np.add.at(weights, pos_sample, 1)
|
|
weights[pos_sample] = 1
|
|
weights[neg_sample] = 1
|
|
# g.edata['sample_weights'] = mx.nd.array(weights, ctx=g.edata['rel_class'].context)
|
|
# return g
|
|
eids = np.where(weights > 0)[0]
|
|
sub_g = g.edge_subgraph(toindex(eids.tolist()))
|
|
sub_g.copy_from_parent()
|
|
sub_g.edata["sample_weights"] = mx.nd.array(
|
|
weights[eids], ctx=g.edata["rel_class"].context
|
|
)
|
|
return sub_g
|