INNER CODE UNIT · Python

step

onnxsim/onnxsim · onnxsim/rpc/server.py:178

        def step(*args):
            return tuple(o.realize() for o in self._forward(dict(zip(names, args))))

        jit = TinyJit(step)
        # Call 1 warms up, call 2 captures, call 3 is the first replay. The capture compiles the
        # graph when the call returns, i.e. outside ``_forward``, so codegen options (BEAM) must
        # be in scope around the JIT calls themselves.
        with self._scope():
            for _ in range(3):
                jit(*tensors.values())
        self._Device[self.device].synchronize()
        results = []
        for _ in range(repeat):
            start = time.perf_counter()
            for _ in range(number):
                jit(*tensors.values())
            self._Device[self.device].synchronize()
            results.append((time.perf_counter() - start) / number)

View source record →