How can I speed up calling python function?


#1

I would like to use the openai gym in julia. Thanks to PyCall this is very easy. The only downside is that my naive approach seems to be pretty slow.

To demonstrate this, I use a minimal example with numpy. I observe qualitatively the same with openai gym (see below).

julia> using PyCall, BenchmarkTools

julia> @pyimport numpy

julia> @btime numpy.tanh(2.)
  6.716 μs (23 allocations: 608 bytes)
0.9640275800758169

julia> @btime pycall(numpy.tanh, Float64, 2.)
  3.409 μs (16 allocations: 400 bytes)
0.9640275800758169

julia> py"""                                                                                   
       import timeit                                                                           
       print(timeit.timeit('numpy.tanh(2.)', setup = 'import numpy', number = 1000000))        
       """
0.5979069629975129

julia> Pkg.status("PyCall")
 - PyCall                        1.15.0

Using pycall improves performance a bit. But it still takes more than 5x more time than the approximately 0.6 micro seconds of calling the function in python.

Here is an example with open ai gym

julia> @pyimport gym.envs.classic_control.cartpole as cartpole
       env = cartpole.CartPoleEnv()
       env[:reset]();
WARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.

julia> @btime env[:step](0)
WARN: You are calling 'step()' even though this environment has already returned done = True. You should always call 'reset()' once you receive 'done = True' -- any further steps are undefined behavior.
  318.556 μs (144 allocations: 5.91 KiB)
([-3.48902e5, -2515.92, 7975.14, 35.7326], 0.0, true, Dict{Any,Any}())

julia> @btime pycall(env[:step], Tuple{Array{Float64, 1}, Float64, Bool, Dict{Any,Any}}, 0)
  261.290 μs (85 allocations: 3.38 KiB)
([-3.00522e6, -7382.4, 27513.3, 35.4553], 0.0, true, Dict{Any,Any}())

julia> py"import timeit
       print(timeit.timeit('env.step(0)', setup = 'from gym.envs.classic_control import cartpole; env = cartpole.CartPoleEnv(); env.reset()', number = 1000000))
       "
WARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.
WARN: You are calling 'step()' even though this environment has already returned done = True. You should always call 'reset()' once you receive 'done = True' -- any further steps are undefined behavior.
4.360256034000486

Here, the version with pycall is approximately 60 times slower than the pure python call.


#2

doesn’t fix the speed problem, but there already a few Gym’s wrappers around, e.g. https://github.com/ozanarkancan/Gym.jl (now registered on metadata)


#3

As I recall there is an enormous amount of overhead involved with actually passing values back and forth between Python and Julia. In theory I suppose it would be possible to use C to circumvent this for things like numpy, but it would be a lot of work and certainly not worth the effort.

Anyway, I know there were (perhaps very hackish) ways of avoiding fetching return values, so in cases where you want to call a function but don’t want the return value, this might be something you can look into (as I recall passing floats and ints to python wasn’t so bad). You might have to delve into the PyCall source code a little bit.