torchsde vs DifferentialEquations.jl / DiffEqFlux.jl (Julia)
This example is a 4-dimensional geometric brownian motion. The code
for the torchsde version is pulled directly from the
so that it would be a fair comparison against the author's own code.
The only change to that example is the addition of a
dt choice so that
the simulation method and time step matches between the two different programs.
The SDE is solved 100 times. The summary of the results is as follows:
- torchsde: 1.87 seconds
- DifferentialEquations.jl: 0.00115 seconds
This demonstrates a 1,600x performance difference in favor of Julia on the Python library's README example. Further testing against torchsde was not able to be completed because of these performance issues.
Note about regularity
We note that the performance difference in the context of neural SDEs is likely smaller due to the ability to time spent in matrix multiplication kernels. However, given that full SDE training examples like demonstrated here generally take about a minute, we still highly expect a major performance difference but currently do not have the compute time to run a full demonstration.
Here's the code using the above described features
It completes in ~5 seconds on my local machine. There is still a difference, but the difference isn't as huge as you described. Do expect further improvements in the near future.
We also want to stress that we are currently more focused on real applications that might be "matmul-heavy" and intended to run on GPUs. But we're also open to suggestions on how to make these tiny CPU-only jobs run faster for torchsde.