Skip to content

Commit

Permalink
Optimize neural net example by adding parallelization to pixel proces…
Browse files Browse the repository at this point in the history
…sing and removing unused Enzyme configuration
  • Loading branch information
dominic-chang committed Jan 28, 2025
1 parent 347396b commit 75e1a3e
Showing 1 changed file with 1 addition and 2 deletions.
3 changes: 1 addition & 2 deletions examples/neural-net-example.jl
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ function (m::ImageModel)(x, ps, st)
emission_vals = zeros(Float64, 1, sze*sze)
for n in 0:1
for i in 1:sze
for j in 1:sze
Threads.@threads for j in 1:sze
pix = pixels[i+(j-1)*sze]
α, β = Krang.screen_coordinate(pix)
T = typeof(α)
Expand Down Expand Up @@ -115,7 +115,6 @@ using Optimization
using OptimizationOptimisers
using StatsBase
using ComponentArrays
Enzyme.Compiler.RunAttributor[] = false

function mse(img1::Matrix{T}, img2::Matrix{T}) where T
mean(((img1 ./ sum(img1)) .- (img2 ./ sum(img2))) .^ 2)
Expand Down

0 comments on commit 75e1a3e

Please sign in to comment.