def push(self, name, server_id, id_tensor, data_tensor): """Push sparse message to target KVServer. Note that push() is an async operation that will return immediately after calling. Parameters ---------- name : str data name server_id : int target server id id_tensor : tensor (mx.ndarray or torch.tensor) a vector storing the ID list data_tensor : tensor (mx.ndarray or torch.tensor) a tensor with the same row size of id """ assert server_id >= 0, 'server_id (%d) cannot be a negative number' % server_id assert server_id < self._server_count, 'server_id (%d) must be smaller than server_count' % server_id assert F.ndim(id_tensor) == 1, 'ID must be a vector.' assert F.shape(id_tensor)[0] == F.shape( data_tensor)[0], 'The data must has the same row size with ID.' msg = KVStoreMsg(type=KVMsgType.PUSH, rank=self._client_id, name=name, id=id_tensor, data=data_tensor) _send_kv_msg(self._sender, msg, server_id)
def __call__(self, edges): sdata = edges.src[self.src_field] edata = edges.data[self.edge_field] # Due to the different broadcasting semantics of different backends, # we need to broadcast the sdata and edata to be of the same rank. rank = max(F.ndim(sdata), F.ndim(edata)) sshape = F.shape(sdata) eshape = F.shape(edata) sdata = F.reshape(sdata, sshape + (1, ) * (rank - F.ndim(sdata))) edata = F.reshape(edata, eshape + (1, ) * (rank - F.ndim(edata))) ret = self.mul_op(sdata, edata) return {self.out_field: ret}
def push(self, name, id_tensor, data_tensor): """Push message to KVServer. Note that push() is an async operation that will return immediately after calling. Parameters ---------- name : str data name id_tensor : tensor (mx.ndarray or torch.tensor) a vector storing the global data ID data_tensor : tensor (mx.ndarray or torch.tensor) a tensor with the same row size of data ID """ assert len(name) > 0, 'name cannot be empty.' assert F.ndim(id_tensor) == 1, 'ID must be a vector.' assert F.shape(id_tensor)[0] == F.shape( data_tensor)[0], 'The data must has the same row size with ID.' # partition data (we can move this part of code into C-api if needed) server_id = self._data_store[name + '-part-'][id_tensor] # sort index by server id sorted_id = F.tensor(np.argsort(F.asnumpy(server_id))) id_tensor = id_tensor[sorted_id] data_tensor = data_tensor[sorted_id] server, count = np.unique(F.asnumpy(server_id), return_counts=True) # push data to server by order start = 0 for idx in range(len(server)): end = start + count[idx] if start == end: # don't have any data for target server continue partial_id = id_tensor[start:end] partial_data = data_tensor[start:end] if server[ idx] in self._local_server_id and self._close_shared_mem == False: if (name + '-g2l-' in self._has_data) == True: local_id = self._data_store[name + '-g2l-'][partial_id] else: local_id = partial_id self._push_handler(name + '-data-', local_id, data_tensor, self._data_store) else: msg = KVStoreMsg(type=KVMsgType.PUSH, rank=self._client_id, name=name, id=partial_id, data=partial_data) _send_kv_msg(self._sender, msg, server[idx]) start += count[idx]
def _is_spmv_supported_edge_feat(g, field): """Return whether the edge feature shape supports SPMV optimization. Only scalar feature is supported currently. """ feat = g.get_e_repr()[field] shape = F.shape(feat) return len(shape) == 1 or (len(shape) == 2 and shape[1] == 1)
def _generate(self, g, eids, canonical_etype): _, _, vtype = canonical_etype shape = F.shape(eids) dtype = F.dtype(eids) ctx = F.context(eids) shape = (shape[0] * self.k,) src, _ = g.find_edges(eids, etype=canonical_etype) src = F.repeat(src, self.k, 0) dst = F.randint(shape, dtype, ctx, 0, g.number_of_nodes(vtype)) return src, dst
def __call__(self, edges): src_data = edges.src[self.src_field] edata = edges.data[self.edge_field] if F.ndim(edata) == 1: # edge feature is a scalar, unsqueeze dims of len 1 src_dim = F.ndim(src_data) new_eshape = (F.shape(edata)[0], ) + (1, ) * (src_dim - 1) edata = F.reshape(edata, new_eshape) ret = self.mul_op(src_data, edata) return {self.out_field: ret}
def push(self, name, ID, data): """Push sparse message to KVServer The push() API will partition message into different KVServer nodes automatically. Note that we assume the row Ids in ID is in the ascending order. Parameters ---------- name : str data name ID : tensor (mx.ndarray or torch.tensor) a vector storing the global IDs data : tensor (mx.ndarray or torch.tensor) a tensor with the same row size of id """ assert F.ndim(ID) == 1, 'ID must be a vector.' assert F.shape(ID)[0] == F.shape( data)[0], 'The data must has the same row size with ID.' group_size = [0] * self._server_count numpy_id = F.asnumpy(ID) count = math.ceil(self._data_size[name] / self._server_count) server_id = numpy_id / count id_list, id_count = np.unique(server_id, return_counts=True) for idx in range(len(id_list)): group_size[int(id_list[idx])] += id_count[idx] min_idx = 0 max_idx = 0 for idx in range(self._server_count): if group_size[idx] == 0: continue max_idx += group_size[idx] range_id = ID[min_idx:max_idx] range_data = data[min_idx:max_idx] min_idx = max_idx msg = KVStoreMsg(type=KVMsgType.PUSH, rank=self._client_id, name=name, id=range_id, data=range_data) _send_kv_msg(self._sender, msg, idx)
def knn_graphE(x, k, istrain=False): """Transforms the given point set to a directed graph, whose coordinates are given as a matrix. The predecessors of each point are its k-nearest neighbors. If a 3D tensor is given instead, then each row would be transformed into a separate graph. The graphs will be unioned. Parameters ---------- x : Tensor The input tensor. If 2D, each row of ``x`` corresponds to a node. If 3D, a k-NN graph would be constructed for each row. Then the graphs are unioned. k : int The number of neighbors Returns ------- DGLGraph The graph. The node IDs are in the same order as ``x``. """ if F.ndim(x) == 2: x = F.unsqueeze(x, 0) n_samples, n_points, _ = F.shape(x) dist = pairwise_squared_distance(x) if istrain and np.random.rand() > 0.5: k_indices = F.argtopk(dist, round(1.5 * k), 2, descending=False) rand_k = np.random.permutation(round(1.5 * k) - 1)[0:k - 1] + 1 # 0 + random k-1 rand_k = np.append(rand_k, 0) k_indices = k_indices[:, :, rand_k] # add 0 else: k_indices = F.argtopk(dist, k, 2, descending=False) dst = F.copy_to(k_indices, F.cpu()) src = F.zeros_like(dst) + F.reshape(F.arange(0, n_points), (1, -1, 1)) per_sample_offset = F.reshape( F.arange(0, n_samples) * n_points, (-1, 1, 1)) dst += per_sample_offset src += per_sample_offset dst = F.reshape(dst, (-1, )) src = F.reshape(src, (-1, )) adj = sparse.csr_matrix( (F.asnumpy(F.zeros_like(dst) + 1), (F.asnumpy(dst), F.asnumpy(src)))) g = DGLGraph(adj, readonly=True) return g
def _generate(self, g, eids, canonical_etype): _, _, vtype = canonical_etype shape = F.shape(eids) dtype = F.dtype(eids) ctx = F.context(eids) shape = (shape[0] * self.k, ) src, _ = g.find_edges(eids, etype=canonical_etype) src = F.repeat(src, self.k, 0) dst = np.random.choice(np.arange(0, g.number_of_nodes()), shape, replace=True, p=self.p) # dst = F.randint(shape, dtype, ctx, 0, g.number_of_nodes(vtype)) dst = th.tensor(dst, dtype=dtype, device=ctx) return src, dst
def push_all(self, name, data): """Push the whole data to KVServer The push_all() API will partition message into different KVServer nodes automatically. Note that we assume the row Ids in ID is in the ascending order. Parameters ---------- name : str data name data : tensor (mx.ndarray or torch.tensor) data tensor """ ID = F.zerocopy_from_numpy(np.arange(F.shape(data)[0])) self.push(name, ID, data)
def segmented_knn_graph(x, k, segs): """Transforms the given point set to a directed graph, whose coordinates are given as a matrix. The predecessors of each point are its k-nearest neighbors. The matrices are concatenated along the first axis, and are segmented by ``segs``. Each block would be transformed into a separate graph. The graphs will be unioned. Parameters ---------- x : Tensor The input tensor. k : int The number of neighbors segs : iterable of int Number of points of each point set. Must sum up to the number of rows in ``x``. Returns ------- DGLGraph The graph. The node IDs are in the same order as ``x``. """ n_total_points, _ = F.shape(x) offset = np.insert(np.cumsum(segs), 0, 0) h_list = F.split(x, segs, 0) dst = [ F.argtopk(pairwise_squared_distance(h_g), k, 1, descending=False) + offset[i] for i, h_g in enumerate(h_list) ] dst = F.cat(dst, 0) src = F.arange(0, n_total_points).unsqueeze(1).expand(n_total_points, k) dst = F.reshape(dst, (-1, )) src = F.reshape(src, (-1, )) # !!! fix shape adj = sparse.csr_matrix( (F.asnumpy(F.zeros_like(dst) + 1), (F.asnumpy(dst), F.asnumpy(src))), shape=(n_total_points, n_total_points)) g = DGLGraph(adj, readonly=True) return g
def score(self, head, rel, tail, triplet_wise=False): head_emb = self.entity_emb(head) rel_emb = self.relation_emb(rel) tail_emb = self.entity_emb(tail) num_head = F.shape(head)[0] num_rel = F.shape(rel)[0] num_tail = F.shape(tail)[0] batch_size = self.batch_size score = [] if triplet_wise: class FakeEdge(object): def __init__(self, head_emb, rel_emb, tail_emb): self._hobj = {} self._robj = {} self._tobj = {} self._hobj['emb'] = head_emb self._robj['emb'] = rel_emb self._tobj['emb'] = tail_emb @property def src(self): return self._hobj @property def dst(self): return self._tobj @property def data(self): return self._robj for i in range((num_head + batch_size - 1) // batch_size): sh_emb = head_emb[i * batch_size : (i + 1) * batch_size \ if (i + 1) * batch_size < num_head \ else num_head] sr_emb = rel_emb[i * batch_size : (i + 1) * batch_size \ if (i + 1) * batch_size < num_head \ else num_head] st_emb = tail_emb[i * batch_size : (i + 1) * batch_size \ if (i + 1) * batch_size < num_head \ else num_head] edata = FakeEdge(sh_emb, sr_emb, st_emb) score.append( F.copy_to( self.score_func.edge_func(edata)['score'], F.cpu())) score = F.cat(score, dim=0) return score else: for i in range((num_head + batch_size - 1) // batch_size): sh_emb = head_emb[i * batch_size : (i + 1) * batch_size \ if (i + 1) * batch_size < num_head \ else num_head] s_score = [] for j in range((num_tail + batch_size - 1) // batch_size): st_emb = tail_emb[j * batch_size : (j + 1) * batch_size \ if (j + 1) * batch_size < num_tail \ else num_tail] s_score.append( F.copy_to( self.score_func.infer(sh_emb, rel_emb, st_emb), F.cpu())) score.append(F.cat(s_score, dim=2)) score = F.cat(score, dim=0) return F.reshape(score, (num_head * num_rel * num_tail, ))
def topK(self, head=None, rel=None, tail=None, exec_mode='all', k=10): if head is None: head = F.arange(0, self.model.num_entity) else: head = F.tensor(head) if rel is None: rel = F.arange(0, self.model.num_rel) else: rel = F.tensor(rel) if tail is None: tail = F.arange(0, self.model.num_entity) else: tail = F.tensor(tail) num_head = F.shape(head)[0] num_rel = F.shape(rel)[0] num_tail = F.shape(tail)[0] if exec_mode == 'triplet_wise': result = [] assert num_head == num_rel, \ 'For triplet wise exection mode, head, relation and tail lists should have same length' assert num_head == num_tail, \ 'For triplet wise exection mode, head, relation and tail lists should have same length' raw_score = self.model.score(head, rel, tail, triplet_wise=True) score = self.score_func(raw_score) idx = F.arange(0, num_head) sidx = F.argsort(score, dim=0, descending=True) sidx = sidx[:k] score = score[sidx] idx = idx[sidx] result.append((F.asnumpy(head[idx]), F.asnumpy(rel[idx]), F.asnumpy(tail[idx]), F.asnumpy(score))) elif exec_mode == 'all': result = [] raw_score = self.model.score(head, rel, tail) score = self.score_func(raw_score) idx = F.arange(0, num_head * num_rel * num_tail) sidx = F.argsort(score, dim=0, descending=True) sidx = sidx[:k] score = score[sidx] idx = idx[sidx] tail_idx = idx % num_tail idx = floor_divide(idx, num_tail) rel_idx = idx % num_rel idx = floor_divide(idx, num_rel) head_idx = idx % num_head result.append((F.asnumpy(head[head_idx]), F.asnumpy(rel[rel_idx]), F.asnumpy(tail[tail_idx]), F.asnumpy(score))) elif exec_mode == 'batch_head': result = [] for i in range(num_head): raw_score = self.model.score(F.unsqueeze(head[i], 0), rel, tail) score = self.score_func(raw_score) idx = F.arange(0, num_rel * num_tail) sidx = F.argsort(score, dim=0, descending=True) sidx = sidx[:k] score = score[sidx] idx = idx[sidx] tail_idx = idx % num_tail idx = floor_divide(idx, num_tail) rel_idx = idx % num_rel result.append((np.full((k,), F.asnumpy(head[i])), F.asnumpy(rel[rel_idx]), F.asnumpy(tail[tail_idx]), F.asnumpy(score))) elif exec_mode == 'batch_rel': result = [] for i in range(num_rel): raw_score = self.model.score(head, F.unsqueeze(rel[i], 0), tail) score = self.score_func(raw_score) idx = F.arange(0, num_head * num_tail) sidx = F.argsort(score, dim=0, descending=True) sidx = sidx[:k] score = score[sidx] idx = idx[sidx] tail_idx = idx % num_tail idx = floor_divide(idx, num_tail) head_idx = idx % num_head result.append((F.asnumpy(head[head_idx]), np.full((k,), F.asnumpy(rel[i])), F.asnumpy(tail[tail_idx]), F.asnumpy(score))) elif exec_mode == 'batch_tail': result = [] for i in range(num_tail): raw_score = self.model.score(head, rel, F.unsqueeze(tail[i], 0)) score = self.score_func(raw_score) idx = F.arange(0, num_head * num_rel) sidx = F.argsort(score, dim=0, descending=True) sidx = sidx[:k] score = score[sidx] idx = idx[sidx] rel_idx = idx % num_rel idx = floor_divide(idx, num_rel) head_idx = idx % num_head result.append((F.asnumpy(head[head_idx]), F.asnumpy(rel[rel_idx]), np.full((k,), F.asnumpy(tail[i])), F.asnumpy(score))) else: assert False, 'unknow execution mode type {}'.format(exec_mode) return result
def start(self): """Start service of KVServer """ # Get connected with all client nodes server_ip, server_port = self._addr.split(':') _receiver_wait(self._receiver, server_ip, int(server_port), self._client_count) # recv client addr and assign ID for clients addr_list = [] for i in range(self._client_count): msg = _recv_kv_msg(self._receiver) assert msg.type == KVMsgType.IP_ID addr_list.append(msg.name) self._sort_addr(addr_list) for ID in range(len(addr_list)): self._client_namebook[ID] = addr_list[ID] _network_wait() for ID, addr in self._client_namebook.items(): client_ip, client_port = addr.split(':') _add_receiver_addr(self._sender, client_ip, int(client_port), ID) _sender_connect(self._sender) if self._server_id == 0: # assign ID to client nodes for client_id, addr in self._client_namebook.items(): msg = KVStoreMsg(type=KVMsgType.IP_ID, rank=self._server_id, name=str(client_id), id=None, data=None) _send_kv_msg(self._sender, msg, client_id) # send serilaized shared-memory tensor information to clients shared_tensor = '' for name in self._has_data: shared_tensor += self._serialize_shared_tensor( name, F.shape(self._data_store[name]), F.dtype(self._data_store[name])) shared_tensor += '|' msg = KVStoreMsg(type=KVMsgType.IP_ID, rank=self._server_id, name=shared_tensor, id=None, data=None) for client_id in range(len(self._client_namebook)): _send_kv_msg(self._sender, msg, client_id) # Service loop while True: msg = _recv_kv_msg(self._receiver) # PUSH message if msg.type == KVMsgType.PUSH: if (msg.name + '-g2l-' in self._has_data) == True: local_id = self._data_store[msg.name + '-g2l-'][msg.id] else: local_id = msg.id self._push_handler(msg.name + '-data-', local_id, msg.data, self._data_store) # PULL message elif msg.type == KVMsgType.PULL: if (msg.name + '-g2l-' in self._has_data) == True: local_id = self._data_store[msg.name + '-g2l-'][msg.id] else: local_id = msg.id res_tensor = self._pull_handler(msg.name + '-data-', local_id, self._data_store) back_msg = KVStoreMsg(type=KVMsgType.PULL_BACK, rank=self._server_id, name=msg.name, id=msg.id, data=res_tensor) _send_kv_msg(self._sender, back_msg, msg.rank) # Barrier message elif msg.type == KVMsgType.BARRIER: self._barrier_count += 1 if self._barrier_count == self._client_count: back_msg = KVStoreMsg(type=KVMsgType.BARRIER, rank=self._server_id, name=None, id=None, data=None) for i in range(self._client_count): _send_kv_msg(self._sender, back_msg, i) self._barrier_count = 0 # FINAL message elif msg.type == KVMsgType.FINAL: print("Exit KVStore service, server ID: %d" % self._server_id) break # exit loop else: raise RuntimeError('Unknown type of kvstore message: %d' % msg.type.value)