tony/tensorneat-mend
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

reduce the use of device_get in speciate

Browse Source
This commit is contained in:
wls2002
2023-05-08 15:58:09 +08:00
parent 91206c796f
commit dde338696f
2 changed files with 12 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
algorithms/neat/species.py
View File

@@ -92,10 +92,10 @@ class SpeciesController:
# First, fast match the population to previous species # First, fast match the population to previous species
if previous_species_list: # exist previous species if previous_species_list: # exist previous species
rid_list = [new_representatives[sid] for sid in previous_species_list] rid_list = [new_representatives[sid] for sid in previous_species_list]
res_pop_distance = [ res_pop_distance = jax.device_get([
jax.device_get(o2m_distance(pop_nodes[rid], pop_connections[rid], pop_nodes, pop_connections)) o2m_distance(pop_nodes[rid], pop_connections[rid], pop_nodes, pop_connections)
for rid in rid_list for rid in rid_list
] ])
pop_res_distance = np.stack(res_pop_distance, axis=0).T pop_res_distance = np.stack(res_pop_distance, axis=0).T
for i in range(pop_res_distance.shape[0]): for i in range(pop_res_distance.shape[0]):
@@ -118,10 +118,10 @@ class SpeciesController:
if len(new_representatives) != 0: if len(new_representatives) != 0:
# the representatives of new species # the representatives of new species
sid, rid = list(zip(*[(k, v) for k, v in new_representatives.items()])) sid, rid = list(zip(*[(k, v) for k, v in new_representatives.items()]))
distances = [ distances = jax.device_get([
jax.device_get(o2o_distance(pop_nodes[i], pop_connections[i], pop_nodes[r], pop_connections[r])) o2o_distance(pop_nodes[i], pop_connections[i], pop_nodes[r], pop_connections[r])
for r in rid for r in rid
] ])
distances = np.array(distances) distances = np.array(distances)
min_idx = np.argmin(distances) min_idx = np.argmin(distances)
min_val = distances[min_idx] min_val = distances[min_idx]

54
examples/jax_playground.py
View File

@@ -8,56 +8,14 @@ from functools import partial
from examples.time_utils import using_cprofile from examples.time_utils import using_cprofile
def func(x, y):
"""
:param x: (100, )
:param y: (100,
:return:
"""
return x * y
def func2(x, y, s):
"""
:param x: (100, )
:param y: (100,
:return:
"""
if s == '123':
return 0
else:
return x * y
@jit @jit
def func3(x, y): def func(x, y):
return func2(x, y, '123') return x + y
# @using_cprofile a, b, c = jnp.array([1]), jnp.array([2]), jnp.array([3])
def main(): li = [a, b, c]
key = jax.random.PRNGKey(42)
x1, y1 = jax.random.normal(key, shape=(1000,)), jax.random.normal(key, shape=(1000,)) cpu_li = jax.device_get(li)
jit_lower_func = jit(func).lower(1, 2).compile() print(cpu_li)
print(type(jit_lower_func))
print(jit_lower_func.memory_analysis())
jit_compiled_func2 = jit(func2, static_argnames=['s']).lower(x1, y1, '123').compile()
print(jit_compiled_func2(x1, y1))
# print(jit_compiled_func2(x1, y1))
f = func3.lower(x1, y1).compile()
print(f(x1, y1))
# print(jit_lower_func(x1, y1))
# x2, y2 = jax.random.normal(key, shape=(200,)), jax.random.normal(key, shape=(200,))
# print(jit_lower_func(x2, y2))
if __name__ == '__main__':
main()