def __init__(self, arg):
super(SkipGram, self).__init__()
self.neg_loss = neg.NEG_loss(num_classes=arg['num_nodes'],
embed_size=arg['emb_size'])
self.SGD = optim.SGD(self.neg_loss.parameters(), lr=0.025)
self.walks = arg['walks']
self.window_size = arg['window_size']
self.walk_length = len(self.walks[0])
self.iter = arg['iter']
self.neg_ratio = arg['neg_ratio']
print self.walk_length
def __init__(self, arg):
super(SkipGram, self).__init__()
type_offset = cPickle.load(
open(arg['data_dir'] + arg['graph_name'] + '_offset.p'))
self.neg_loss = neg.NEG_loss(type_offset=type_offset,
node_types=arg['node_types'],
edge_types=arg['edge_types'],
embed_size=arg['emb_size'],
pre_train_path=arg['pre_train'],
graph_name=arg['graph_name'],
mode=arg['mode'],
map_mode=arg['map_mode'])
self.input = arg['network']
#print('edge layer learning rate is:', arg['lr'] * (float(len(arg['edge_types'])) / type_offset['sum']))
edge_stats = cPickle.load(
open(arg['data_dir'] + arg['graph_name'] + '_edge_stat.p'))
#print(edge_stats)
self.mode = arg['mode']
self.map_mode = arg['map_mode']
self.dump_timer = arg['dump_timer']
self.model_dir = arg['model_dir']
self.log_dir = arg['log_dir']
self.more_param = arg['more_param']
self.fine_tune = arg['fine_tune']
self.lr = arg['lr']
self.weighted = arg['weighted']
self._params = []
if self.map_mode != -1:
self._params = [{
'params': self.neg_loss.in_embed.parameters()
}, {
'params': self.neg_loss.out_embed.parameters()
}]
for i in xrange(len(self.neg_loss.edge_mapping)):
self._params.append({
'params':
self.neg_loss.edge_mapping[i].parameters(),
'lr':
arg['lr'] * arg['lr_ratio'] * (float(len(self.input))) /
(type_offset['sum'] * edge_stats[i] + 1e-6)
})
self.window_size = arg['window_size']
self.graph_name = arg['graph_name']
self.data = tdata.DataLoader(self.input,
arg['batch_size'],
shuffle=True)
self.batch_size = arg['batch_size']
self.iter = arg['iter']
self.neg_ratio = arg['neg_ratio']
args.dimensions, args.graph_name,
config['nodes'])
#args.op = 1
#print(_data)
t.cuda.set_device(int(args.gpu))
type_offset = cPickle.load(
open(args.data_dir + args.graph_name + '_offset.p'))
in_mapping = cPickle.load(
open(args.data_dir + args.graph_name + '_in_mapping.p'))
out_mapping = cPickle.load(
open(args.data_dir + args.graph_name + '_out_mapping.p'))
model = neg.NEG_loss(type_offset=type_offset,
node_types=config['nodes'],
edge_types=config['edges'],
embed_size=args.dimensions,
pre_train_path=_data,
graph_name=args.graph_name,
mode=args.op,
map_mode=args.map_func)
#print(model.in_embed.weight.sum())
if args.op != -1:
if args.more_param != 'None':
model_path = args.model_dir + 'heer_' + args.graph_name + '_' + str(args.iter) + '_op_' + str(args.op) + \
'_mode_' + str(args.map_func)+ '_' + args.more_param + '.pt'
else:
model_path = args.model_dir + 'heer_' + args.graph_name + '_' + str(args.iter) + '_op_' + str(args.op) + \
'_mode_' + str(args.map_func)+ '.pt'
print('model path:', model_path)
xxx = t.load(model_path, map_location=lambda storage, loc: storage)
model.load_state_dict(xxx, False)
