Small change to testspeed file.
PiperOrigin-RevId: 618323795 Change-Id: I3a6307b8ae6d9af0d17c0ae85eae673d60112ddf
This commit is contained in:
Baruch Tabanpour
committed by
Copybara-Service
Copybara-Service
parent
4660a297d1
commit
2019d42ca6
+41
-84
@@ -14,110 +14,67 @@
|
||||
# ==============================================================================
|
||||
"""Run benchmarks on various devices."""
|
||||
|
||||
import os
|
||||
import time
|
||||
from typing import Sequence, Tuple
|
||||
from typing import Sequence
|
||||
|
||||
from absl import app
|
||||
from absl import flags
|
||||
from etils import epath
|
||||
import jax
|
||||
import mujoco
|
||||
from mujoco import mjx
|
||||
|
||||
FLAGS = flags.FLAGS
|
||||
|
||||
flags.DEFINE_string('mjcf', None, 'path to model', required=True)
|
||||
flags.DEFINE_integer('nstep', 1000, 'number of steps per rollout')
|
||||
flags.DEFINE_integer('batch_size', 1024, 'number of parallel rollouts')
|
||||
flags.DEFINE_integer('unroll', 1, 'loop unroll length')
|
||||
flags.DEFINE_enum('solver', 'cg', ['cg', 'newton'], 'constraint solver')
|
||||
flags.DEFINE_integer('iterations', 1, 'number of solver iterations')
|
||||
flags.DEFINE_integer('ls_iterations', 4, 'number of linesearch iterations')
|
||||
flags.DEFINE_enum('output', 'text', ['text', 'tsv'], 'format to print results')
|
||||
|
||||
|
||||
def _measure(fn, *args) -> Tuple[float, float]:
|
||||
"""Reports jit time and op time for a function."""
|
||||
|
||||
beg = time.perf_counter()
|
||||
compiled_fn = fn.lower(*args).compile()
|
||||
end = time.perf_counter()
|
||||
jit_time = end - beg
|
||||
|
||||
beg = time.perf_counter()
|
||||
result = compiled_fn(*args)
|
||||
jax.block_until_ready(result)
|
||||
end = time.perf_counter()
|
||||
run_time = end - beg
|
||||
|
||||
return jit_time, run_time
|
||||
_MJCF = flags.DEFINE_string('mjcf', None, 'path to model', required=True)
|
||||
_BASE_PATH = flags.DEFINE_string(
|
||||
'base_path', None, 'base path, defaults to mujoco.mjx resource path'
|
||||
)
|
||||
_NSTEP = flags.DEFINE_integer('nstep', 1000, 'number of steps per rollout')
|
||||
_BATCH_SIZE = flags.DEFINE_integer(
|
||||
'batch_size', 1024, 'number of parallel rollouts'
|
||||
)
|
||||
_UNROLL = flags.DEFINE_integer('unroll', 1, 'loop unroll length')
|
||||
_SOLVER = flags.DEFINE_enum(
|
||||
'solver', 'cg', ['cg', 'newton'], 'constraint solver'
|
||||
)
|
||||
_ITERATIONS = flags.DEFINE_integer(
|
||||
'iterations', 1, 'number of solver iterations'
|
||||
)
|
||||
_LS_ITERATIONS = flags.DEFINE_integer(
|
||||
'ls_iterations', 4, 'number of linesearch iterations'
|
||||
)
|
||||
_OUTPUT = flags.DEFINE_enum(
|
||||
'output', 'text', ['text', 'tsv'], 'format to print results'
|
||||
)
|
||||
|
||||
|
||||
def _main(argv: Sequence[str]):
|
||||
"""Benchmark a model."""
|
||||
|
||||
xla_flags = os.environ.get('XLA_FLAGS', '')
|
||||
xla_flags += ' --xla_gpu_triton_gemm_any=True'
|
||||
os.environ['XLA_FLAGS'] = xla_flags
|
||||
|
||||
f = epath.resource_path('mujoco.mjx') / 'test_data' / FLAGS.mjcf
|
||||
"""Runs testpeed function."""
|
||||
base_path = _BASE_PATH.value or epath.resource_path('mujoco.mjx')
|
||||
f = base_path / 'test_data' / _MJCF.value
|
||||
m = mujoco.MjModel.from_xml_path(f.as_posix())
|
||||
m.opt.solver = {
|
||||
'cg': mujoco.mjtSolver.mjSOL_CG,
|
||||
'newton': mujoco.mjtSolver.mjSOL_NEWTON,
|
||||
}[FLAGS.solver.lower()]
|
||||
m.opt.iterations = FLAGS.iterations
|
||||
m.opt.ls_iterations = FLAGS.ls_iterations
|
||||
m = mjx.put_model(m)
|
||||
|
||||
if FLAGS.output == 'text':
|
||||
print(f"Rolling out {FLAGS.nstep} steps at dt = {m.opt.timestep:.3f}...")
|
||||
print(f'Rolling out {_NSTEP.value} steps at dt = {m.opt.timestep:.3f}...')
|
||||
jit_time, run_time, steps = mjx.benchmark(
|
||||
m,
|
||||
_NSTEP.value,
|
||||
_BATCH_SIZE.value,
|
||||
_UNROLL.value,
|
||||
_SOLVER.value,
|
||||
_ITERATIONS.value,
|
||||
_LS_ITERATIONS.value,
|
||||
)
|
||||
|
||||
@jax.pmap
|
||||
def init(key):
|
||||
key = jax.random.split(key, FLAGS.batch_size // jax.device_count())
|
||||
|
||||
@jax.vmap
|
||||
def random_init(key):
|
||||
d = mjx.make_data(m)
|
||||
qvel = 0.01 * jax.random.normal(key, shape=(m.nv,))
|
||||
d = d.replace(qvel=qvel)
|
||||
return d
|
||||
|
||||
return random_init(key)
|
||||
|
||||
key = jax.random.split(jax.random.key(0), jax.device_count())
|
||||
d = init(key)
|
||||
jax.block_until_ready(d)
|
||||
|
||||
@jax.pmap
|
||||
def unroll(d):
|
||||
|
||||
@jax.vmap
|
||||
def step(d, _):
|
||||
d = mjx.step(m, d)
|
||||
return d, None
|
||||
|
||||
d, _ = jax.lax.scan(step, d, None, length=FLAGS.nstep, unroll=FLAGS.unroll)
|
||||
|
||||
return d
|
||||
|
||||
jit_time, run_time = _measure(unroll, d)
|
||||
steps = FLAGS.nstep * FLAGS.batch_size
|
||||
|
||||
if FLAGS.output == 'text':
|
||||
name = argv[0]
|
||||
if _OUTPUT.value == 'text':
|
||||
print(f"""
|
||||
Summary for {FLAGS.batch_size} parallel rollouts
|
||||
Summary for {_BATCH_SIZE.value} parallel rollouts
|
||||
|
||||
Total JIT time: {jit_time:.2f} s
|
||||
Total simulation time: {run_time:.2f} s
|
||||
Total steps per second: { steps / run_time:.0f}
|
||||
Total realtime factor: { steps * m.opt.timestep / run_time:.2f} x
|
||||
Total time per step: { 1e6 * run_time / steps:.2f} µs""")
|
||||
elif FLAGS.output == 'tsv':
|
||||
name = argv[0].split('/')[-1].replace('testspeed_', '')
|
||||
print(f"{name}\tjit: {jit_time:.2f}s\tsteps/second: {steps / run_time:.0f}")
|
||||
elif _OUTPUT.value == 'tsv':
|
||||
name = name.split('/')[-1].replace('testspeed_', '')
|
||||
print(f'{name}\tjit: {jit_time:.2f}s\tsteps/second: {steps / run_time:.0f}')
|
||||
|
||||
|
||||
def main():
|
||||
|
||||
Reference in New Issue
Block a user