dmlc--dgl
Tong He
cbee427839
* add working scripts * add frcnn training script * remove redundent files * refactor validation computation, will optimize sgdet and training * validation finally finished * f-rcnn training * test reldn * rm file * update reldn training * data preprocess to h5 * temp * use coco json * fix conflict * new obj dataset for detection * update training * before cleanup * remove abundant files * add arg parse to train * cleanup code file * update * fix * add readme * add ipynb as demo * add demo pic * update readme * add demo script * improve paths * improve readme * add docstrings * fix args description * update readme * add models from s3 * update README Co-authored-by: Minjie Wang <minjie.wang@nyu.edu>
32 行
1.2 KiB
Python
32 行
1.2 KiB
Python
import dgl
|
|
from dgl.utils import toindex
|
|
import mxnet as mx
|
|
import numpy as np
|
|
|
|
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
|