# Converting PyTorch to Flux while keeping performance

**URL:** <https://discourse.julialang.org/t/converting-pytorch-to-flux-while-keeping-performance/81877>\
**Category:** Machine Learning\
**Tags:** flux, zygote\
**Created:** [May 29, 2022, 5:49pm UTC](https://discourse.julialang.org/t/converting-pytorch-to-flux-while-keeping-performance/81877 "2022-05-29T17:49:09Z")\
**Posts on this page:** 1\
**Showing post:** 3

<div class="post-metadata">

**Author:** ![Jarartur](https://avatars.discourse-cdn.com/v4/letter/j/b2d939/32.png) [@Jarartur](https://discourse.julialang.org/u/Jarartur)\
**Post date:** [May 29, 2022, 7:23pm UTC](https://discourse.julialang.org/t/converting-pytorch-to-flux-while-keeping-performance/81877/3 "2022-05-29T19:23:25Z")

</div>

Yes, the loss\_fn is similar to what I get.  
I might just misunderstood you on a previous post but I got that from here: [Zygote Performance (Again…) - General Usage - JuliaLang](https://discourse.julialang.org/t/zygote-performance-again/62753/14)

> Back to the topic at hand, the long and short of it seems to be that there aren’t optimized adjoints in place for mapreduce-like functions (including `sum(f, ...)` ) yet

Now I rewritten the function in Julia using Tullio (also recommended by someone):

```julia
rowmap(f, xs) = @tullio ys[r] := (f)(xs[r,c]) grad=Dual

function stiffness_reg(deformation, idx)

    deformation_indexed = deformation[idx, :]
    deformations1 = (deformation_indexed[:, 1, :] .- deformation_indexed[:, 2, :])
    deformations2 = (deformation_indexed[:, 1, :] .- deformation_indexed[:, 3, :])
    deformations3 = (deformation_indexed[:, 1, :] .- deformation_indexed[:, 4, :])

    norm_vec1 = rowmap(norm, deformations1)
    norm_vec2 = rowmap(norm, deformations2)
    norm_vec3 = rowmap(norm, deformations3)

    sum(norm_vec1+norm_vec2+norm_vec3)
end

```

And the benchmarks are:  
PyTorch

```julia
<torch.utils.benchmark.utils.common.Measurement object at 0x00000133FCB28DC0>
reg(x, y)
setup: from __main__ import reg
  6.68 ms
  1 measurement, 1000 runs , 1 thread

```

Julia

```julia
@benchmark grad = gradient(params) do
           new_points = register(data_src)
           loss = stiffness_reg(register.vec, vec_idxs)
           return loss
       end
BenchmarkTools.Trial: 14 samples with 1 evaluation.
 Range (min … max): 89.977 ms … 1.889 s ┊ GC (min … max): 0.00% … 0.00%
 Time (median): 127.545 ms ┊ GC (median): 0.00%
 Time (mean ± σ): 462.202 ms ± 683.802 ms ┊ GC (mean ± σ): 0.00% ± 0.00%

  █▃            
  ██▄▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▄▁▁▁▁▁▁▁▁▁▁▁▁▄▁▄ ▁
  90 ms Histogram: frequency by time 1.89 s <

 Memory estimate: 101.66 KiB, allocs estimate: 2009.

```

---

_[View the full topic](https://discourse.julialang.org/t/converting-pytorch-to-flux-while-keeping-performance/81877)._