def add_node(self, key):
# print(np.array(key).shape)
if self.gpu_capacity[0] >= self.capacity:
# find the LRU entry
old_index = int(np.argmin(self.lru))
for action in range(self.num_actions):
for successor in self.next_id[old_index][action].keys():
for s, a in self.prev_id[successor]:
if s == old_index:
self.prev_id[successor].remove((s, a))
self.next_id[old_index][action] = dict()
self.states[old_index] = key
self.external_value[old_index] = np.full((self.num_actions,), np.nan)
self.internal_value[old_index] = self.rmax * np.ones(self.num_actions)
self.state_value_u[old_index] = np.nan
self.state_value_v[old_index] = np.nan
self.lru[old_index] = self.tm
self.count[old_index] = 2
self.prev_id[old_index] = []
knn_cuda_fixmem.add(self.address, int(old_index), np.array(key))
self.tm += 0.01
self.newly_added[old_index] = True
return old_index, True
else:
self.states[self.gpu_capacity[0]] = key
self.lru[self.gpu_capacity[0]] = self.tm
self.newly_added[self.gpu_capacity[0]] = True
self.count[self.gpu_capacity[0]] = 2
knn_cuda_fixmem.add(self.address, int(self.gpu_capacity[0]), np.array(key))
self.gpu_capacity[0] += 1
self.tm += 0.01
return self.gpu_capacity[0] - 1, False
def update(self, indexes, z_new):
self.log("update in buffer", self.gpu_capacity[0])
assert len(indexes) == len(z_new), "{}{}".format(len(indexes), len(z_new))
assert z_new.shape[1] == self.z_dim
for i, ind in enumerate(indexes):
self.states[ind] = z_new[i]
knn_cuda_fixmem.add(self.address, int(ind), np.array(z_new[i]).squeeze())
def add_node(self, key):
# print(np.array(key).shape)
if self.curr_capacity >= self.capacity:
# find the LRU entry
old_index = np.argmin(self.lru)
for successor in self.next_id[old_index]:
for s, a in self.prev_id[successor]:
if s == old_index:
self.prev_id.remove((s, a))
self.states[old_index] = key
self.external_value[old_index] = -self.rmax * np.ones(
self.num_actions)
self.internal_value[old_index] = np.zeros(self.num_actions)
self.state_value[old_index] = -self.rmax
self.lru[old_index] = self.tm
self.count[old_index] = 2
self.prev_id = []
knn_cuda_fixmem.add(self.address, old_index, np.array(key))
self.tm += 0.01
return old_index
else:
self.states[self.curr_capacity] = key
self.lru[self.curr_capacity] = self.tm
self.count[self.curr_capacity] = 2
knn_cuda_fixmem.add(self.address, self.curr_capacity,
np.array(key))
self.curr_capacity += 1
self.tm += 0.01
return self.curr_capacity - 1
def add(self, action, key, value, reward, done, brothers):
buffer = self.ec_buffer[action]
if buffer.curr_capacity >= buffer.capacity:
# find the LRU entry
index = np.argmin(buffer.lru)
self.ec_buffer[action].prev_id[index] = []
self.ec_buffer[action].internal_value[index] = 0
for bro in self.ec_buffer[action].brothers[index]:
self.ec_buffer[bro[1]].brothers[bro[0]].remove((index, action))
self.ec_buffer[action].brothers[index] = []
else:
index = buffer.curr_capacity
buffer.curr_capacity += 1
buffer.states[index] = key
for a in range(len(brothers)):
if brothers[a] > 0:
buffer.brothers[index].append((brothers[a], a))
self.ec_buffer[a].brothers[brothers[a]].append((index, action))
buffer.external_value[index] = value
buffer.reward[index] = reward
buffer.done[index] = done
buffer.lru[index] = buffer.tm
buffer.tm += 0.01
# print("here")
# print(action, index, buffer.capacity)
# print(buffer.address)
key = np.array(key, copy=True)
# key = copy.deepcopy(key)
# print(key.shape)
knn_cuda_fixmem.add(buffer.address, index, key)
return index
def add(self, key, external_value, internal_value):
# print(np.array(key).shape)
if self.curr_capacity >= self.capacity:
# find the LRU entry
old_index = np.argmin(self.lru)
self.states[old_index] = key
self.external_value[old_index] = external_value
self.internal_value[old_index] = internal_value
self.lru[old_index] = self.tm
self.count[old_index] = 2
knn_cuda_fixmem.add(self.address, int(old_index), np.array(key))
else:
self.states[self.curr_capacity] = key
self.external_value[self.curr_capacity] = external_value
self.internal_value[self.curr_capacity] = internal_value
self.lru[self.curr_capacity] = self.tm
self.count[self.curr_capacity] = 2
knn_cuda_fixmem.add(self.address, int(self.curr_capacity), np.array(key))
self.curr_capacity += 1
self.tm += 0.01
def add(self, key, value_decay, action=-1):
# print(np.array(key).shape)
if self.curr_capacity >= self.capacity:
# find the LRU entry
old_index = np.argmin(self.lru)
self.states[old_index] = key
self.q_values_decay[old_index] = value_decay
self.lru[old_index] = self.tm
self.count[old_index] = 2
knn_cuda_fixmem.add(self.address, old_index, np.array(key))
if action >= 0:
self.best_action[old_index, action] = 1
else:
self.states[self.curr_capacity] = key
self.q_values_decay[self.curr_capacity] = value_decay
self.lru[self.curr_capacity] = self.tm
self.count[self.curr_capacity] = 2
knn_cuda_fixmem.add(self.address, self.curr_capacity, np.array(key))
if action >= 0:
self.best_action[self.curr_capacity, action] = 1
self.curr_capacity += 1
self.tm += 0.01
def add_node(self, key, knn_dist=[], knn_ind=[]):
# print(np.array(key).shape)
if self.curr_capacity >= self.capacity:
# print("Currently not allow buffer to delete old node!")
# raise NotImplementedError
# find the LRU entry
old_index = int(np.argmin(self.lru))
for action in range(self.num_actions):
for successor in self.next_id[old_index][action].keys():
for s, a in self.prev_id[successor]:
if s == old_index:
self.prev_id[successor].remove((s, a))
self.next_id[old_index][action] = dict()
self.states[old_index] = key
self.external_value[old_index] = np.full((self.num_actions,), np.nan)
self.internal_value[old_index] = self.rmax * np.ones(self.num_actions)
self.state_value_u[old_index] = np.nan
self.state_value_v[old_index] = np.nan
self.lru[old_index] = self.tm
self.count[old_index] = 2
self.prev_id[old_index] = []
knn_cuda_fixmem.add(self.address, int(old_index), np.array(key))
self.tm += 0.01
self.newly_added[old_index] = True
return old_index, True
else:
self.states[self.curr_capacity] = key
self.lru[self.curr_capacity] = self.tm
self.newly_added[self.curr_capacity] = True
self.count[self.curr_capacity] = 2
knn_cuda_fixmem.add(self.address, int(self.curr_capacity), np.array(key))
self.curr_capacity += 1
self.tm += 0.01
return self.curr_capacity - 1, False
def add_node(self, key, knn_dist, knn_ind):
knn_dist = np.array(knn_dist).reshape(-1).tolist()
knn_ind = np.array(knn_ind).reshape(-1).tolist()
# print(np.array(key).shape)
if self.curr_capacity >= self.capacity:
# find the LRU entry
old_index = int(
np.argmin(self.density + 2 * self.rmax * self.newly_added +
self.lru * 0.1))
self.log("density based cache index:", old_index, "density:",
self.density[old_index])
# old_index = np.random.randint(0,self.capacity)
# deal with model
for action in range(self.num_actions):
for successor in self.next_id[old_index][action].keys():
for s, a in self.prev_id[successor]:
if s == old_index:
self.prev_id[successor].remove((s, a))
self.next_id[old_index][action] = dict()
self.prev_id[old_index] = []
# deal with neighbour
for neighbour in self.neighbour_backward[old_index]:
place = np.where(
np.array(self.neighbour_forward[neighbour]) ==
old_index)[0]
if len(place) == 0:
print("backward", old_index,
self.neighbour_backward[old_index])
print("forward", neighbour,
self.neighbour_forward[neighbour])
raise SyntaxError
assert len(place) == 1, place
place = int(place)
self.neighbour_forward[neighbour].pop(place)
self.neighbour_dist_forward[neighbour].pop(place)
# print(self.neighbour_dist_forward[neighbour])
self.compute_density(neighbour)
for neighbour in self.neighbour_forward[old_index]:
self.neighbour_backward[neighbour].remove(old_index)
# clean buffer and then treat it like a new one
self.density[old_index] = self.rmax
self.external_value[old_index] = np.full((self.num_actions, ),
np.nan)
self.internal_value[old_index] = self.rmax * np.ones(
self.num_actions)
self.state_value_u[old_index] = np.nan
self.state_value_v[old_index] = np.nan
self.neighbour_forward[old_index] = []
self.neighbour_dist_forward[old_index] = []
self.neighbour_backward[old_index] = []
new_index = old_index
else:
new_index = self.curr_capacity
self.curr_capacity += 1
if new_index in knn_ind:
place = np.where(np.array(new_index) == knn_ind)[0]
assert len(place) == 1, place
place = int(place)
knn_ind.pop(place)
knn_dist.pop(place)
self.neighbour_forward[new_index] = knn_ind
self.neighbour_dist_forward[new_index] = knn_dist
# print(self.neighbour_dist_forward[new_index])
self.compute_density(new_index)
# print(knn_ind,knn_dist)
# for x in zip([knn_ind, knn_dist]):
# print(x)
for ind, dist in zip(knn_ind, knn_dist):
self.neighbour_backward[ind].append(new_index)
if len(
self.neighbour_forward[ind]
) < self.knn * 4 and new_index not in self.neighbour_forward[ind]:
self.neighbour_forward[ind].append(new_index)
self.neighbour_dist_forward[ind].append(dist)
# print(self.neighbour_dist_forward[ind])
self.compute_density(ind)
self.neighbour_backward[new_index].append(ind)
else:
dist_max = np.max(self.neighbour_dist_forward[ind])
ind_max = int(np.argmax(self.neighbour_dist_forward[ind]))
if dist_max > dist:
ind_tmp = self.neighbour_forward[ind][ind_max]
try:
self.neighbour_backward[ind_tmp].remove(ind)
except ValueError:
print(ind_tmp, self.neighbour_forward[ind])
print(ind, self.neighbour_backward[ind_tmp])
raise ValueError
self.neighbour_forward[ind][ind_max] = new_index
self.neighbour_dist_forward[ind][ind_max] = dist
# print(self.neighbour_dist_forward[ind])
self.compute_density(ind)
self.neighbour_backward[new_index].append(ind)
self.states[new_index] = key
self.lru[new_index] = self.tm
self.newly_added[new_index] = True
self.count[new_index] = 2
knn_cuda_fixmem.add(self.address, int(new_index), np.array(key))
self.tm += 0.01
return new_index, False
k = 4
for n in range(4):
cur_time = time.time()
query = np.random.rand(query_size, c).astype(np.float32)
reference = np.random.rand(dict_size, c).astype(np.float32)
address = knn.allocate(dict_size, c, query_max, k)
# print(address, address.dtype)
# address = copy.deepcopy(address)
# print(address,address.dtype)
# print("??????")
print(address)
for i in range(capacity):
# print(i,np.array(reference[i]).shape)
# print(i)
knn.add(address, i, reference[i])
print("add time:", time.time() - cur_time)
cur_time = time.time()
# print(address)
# # Index is 1-based
dist, ind = knn.knn(address, query.reshape(-1, c), k, capacity)
print(ind.shape, "time:", time.time() - cur_time)
print(np.transpose(ind))
print(np.transpose(dist))
# cur_time = time.time()
tree = KDTree(reference[:capacity])
print("build tree time:", time.time() - cur_time)
cur_time = time.time()
dist, ind = tree.query(query, k=k)