Hello – Is it possible to use automatic differentiation (specifically in Optim,jl) with multi-threading? I’m getting type conflicts with race-safe Atomic type in my main objective functions (a likelihood) not being compatible with the ForwardDiff.Dual type needed for autodiff.
Here is sample code of the objective function:
function ll_nlog2_k2_mtr(b, choicearray, ttarray, ddarray, ffarray, hvec, wtvec)
n = size(choicearray)[1]
kap = b[1:2]
lam = exp.(b[3:9])
lam_noveh = vcat(1.0, lam[2:7])
bshift = 9
mu = b[bshift+1:bshift+16]
T = promote_type(eltype(b))
#ll = zero(T)
ll = Atomic{T}(0)
@threads for i in 1:n
if hvec[i] == 1
pr = ill_nlog2_veh_k2(kap, lam, mu, ttarray[i, :], ddarray[i, :], ffarray[i, :], choicearray[i, 1], choicearray[i, 2], choicearray[i, 3])
elseif hvec[i] == 0
pr = ill_nlog2_noveh_k2(kap, lam_noveh, mu, ttarray[i, :], ddarray[i, :], ffarray[i, :], choicearray[i, 2], choicearray[i, 3])
end
#ll += -1 * wtvec[i] * log(max(eps(T), pr))
atomic_add!(ll, -1 * wtvec[i] * log(max(eps(T), pr)))
end
return ll
end
where the commented sections are those that I altered to create the multi-threaded version. Here is the code that calls it and invokes autodiff:
function estwrap_nlog2_threadtest(dmat)
tt = hcat(Array{Float64}(dmat[:,17:30]), dmat[:,16]) # oss - ott then d
dd = Array{Int64}(dmat[:,32:45])
ff = Array{Float64}(dmat[:,47:60])
choices = Array{Int64}(dmat[:,12:14])
wts = Vector{Float64}(dmat[:,15])
hvec = Vector{Int64}(dmat[:,7])
bni = vcat(-0.05, -0.05, 0.0*ones(7), zeros(14), -0.1, -1.0)
td_newlog = TwiceDifferentiable(vars -> ll_nlog2_k2_mtr(vars, choices, tt, dd, ff, hvec, wts), bni; autodiff = :forward)
b_newlog = optimize(td_newlog, bni, method = NewtonTrustRegion(), iterations = 200, show_trace = true, show_every = 1, g_tol = 1e-4)
bnewlog = Optim.minimizer(b_newlog)
end
I’m able to to run the objective function (evaluate the ll) with multiple threads and am getting a decent speed increase, but to do so in a race-safe way we have to declare ll to be an Atomic type (here Atomic{T} so it inherits the type of b which works fine when I don’t need autodiff.
But with autodiff, I get the type error below (it seems like the only errors I ever get in Julia are type errors, it’s taking years off my life lol):
TypeError: in Atomic, in T, expected T<:Union{Bool, Float16, Float32, Float64, Int128, Int16, Int32, Int64, Int8, UInt128, UInt16, UInt32, UInt64, UInt8}, got Type{ForwardDiff.Dual{ForwardDiff.Tag{var"#38#39"{Vector{Int64}, Vector{Float64}, Matrix{Int64}, Matrix{Float64}, Matrix{Int64}, Matrix{Float64}}, Float64}, Float64, 9}}
Do I need this to somehow be a non-existent combo of atomic and forwarddiff’s dual, or is there a smart (or dumb ha) way to deal with this?
Here is a MWE:
using Optim, Distributions, Random, Statistics, BenchmarkTools, NLSolversBase, LineSearches, ForwardDiff
using Base.Threads
import Base.Threads.@threads
dr = rand(Normal(0.1, 1.1), (30,2))
dta = 1.0 .+ 0.3*dr[:,1] + dr[:,2]
# no thread
function minss(b, y, x)
n = size(y)[1]
T = promote_type(eltype(b))
ss = zero(T)
for i in 1:n
res = (y[i] - b[1] - b[2]*x[i])^2
ss += res
end
return ss
end
function estwrap_minss(y, x)
bni = vcat(0.0, 0.0)
tdf = TwiceDifferentiable(vars -> minss(vars, y, x), bni; autodiff = :forward)
minout = optimize(tdf, bni, method = NewtonTrustRegion(), iterations = 2000, show_trace = true, show_every = 10, g_tol = 1e-4)
b = Optim.minimizer(minout)
end
# with threading
function minss_thr(b, y, x)
n = size(y)[1]
T = promote_type(eltype(b))
ss = Atomic{T}(0)
@threads for i in 1:n
res = (y[i] - b[1] - b[2]*x[i])^2
atomic_add!(ss, res)
end
return ss
end
function estwrap_minss_thr(y, x)
bni = vcat(0.0, 0.0)
tdf = TwiceDifferentiable(vars -> minss_thr(vars, y, x), bni; autodiff = :forward)
minout = optimize(tdf, bni, method = NewtonTrustRegion(), iterations = 2000, show_trace = true, show_every = 10, g_tol = 1e-4)
b = Optim.minimizer(minout)
end
aa1 = estwrap_minss(dta, dr[:,1])
aa2 = estwrap_minss_thr(dta, dr[:,1])