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
@@ -49,4 +49,5 @@ from mujoco.mjx._src.solver import solve
|
||||
from mujoco.mjx._src.support import full_m
|
||||
from mujoco.mjx._src.support import is_sparse
|
||||
from mujoco.mjx._src.support import mul_m
|
||||
from mujoco.mjx._src.test_util import benchmark
|
||||
from mujoco.mjx._src.types import *
|
||||
|
||||
@@ -13,7 +13,6 @@
|
||||
# limitations under the License.
|
||||
# ==============================================================================
|
||||
"""Engine support functions."""
|
||||
|
||||
from typing import Optional, Tuple, Union
|
||||
|
||||
import jax
|
||||
|
||||
@@ -14,14 +14,96 @@
|
||||
# ==============================================================================
|
||||
"""Utilities for testing."""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from typing import Dict, Tuple
|
||||
from xml.etree import ElementTree as ET
|
||||
|
||||
from etils import epath
|
||||
import jax
|
||||
import mujoco
|
||||
# pylint: disable=g-importing-member
|
||||
from mujoco.mjx._src import forward
|
||||
from mujoco.mjx._src import io
|
||||
# pylint: enable=g-importing-member
|
||||
import numpy as np
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
def benchmark(
|
||||
m: mujoco.MjModel,
|
||||
nstep: int = 1000,
|
||||
batch_size: int = 1024,
|
||||
unroll_steps: int = 1,
|
||||
solver: str = 'cg',
|
||||
iterations: int = 1,
|
||||
ls_iterations: int = 4,
|
||||
) -> Tuple[float, float, int]:
|
||||
"""Benchmark a model."""
|
||||
|
||||
xla_flags = os.environ.get('XLA_FLAGS', '')
|
||||
xla_flags += ' --xla_gpu_triton_gemm_any=True'
|
||||
os.environ['XLA_FLAGS'] = xla_flags
|
||||
|
||||
m.opt.solver = {
|
||||
'cg': mujoco.mjtSolver.mjSOL_CG,
|
||||
'newton': mujoco.mjtSolver.mjSOL_NEWTON,
|
||||
}[solver.lower()]
|
||||
m.opt.iterations = iterations
|
||||
m.opt.ls_iterations = ls_iterations
|
||||
m = io.put_model(m)
|
||||
|
||||
@jax.pmap
|
||||
def init(key):
|
||||
key = jax.random.split(key, batch_size // jax.device_count())
|
||||
|
||||
@jax.vmap
|
||||
def random_init(key):
|
||||
d = io.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 = forward.step(m, d)
|
||||
return d, None
|
||||
|
||||
d, _ = jax.lax.scan(step, d, None, length=nstep, unroll=unroll_steps)
|
||||
|
||||
return d
|
||||
|
||||
jit_time, run_time = _measure(unroll, d)
|
||||
steps = nstep * batch_size
|
||||
|
||||
return jit_time, run_time, steps
|
||||
|
||||
|
||||
_ACTUATOR_TYPES = ['motor', 'velocity', 'position', 'general', 'intvelocity']
|
||||
_DYN_TYPES = ['none', 'integrator', 'filter', 'filterexact']
|
||||
_DYN_PRMS = ['0.189', '2.1']
|
||||
|
||||
+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