Closed newalexander closed 2 years ago
Thanks for this very clear report! The issue could be in Chainrule.ProjectTo{SparseMatrixCSC}
not being ForwardDiff.Dual friendly.
It works fine for the generic message passing framework (e.g. for th GATConv)
julia> using GraphNeuralNetworks, Flux, Graphs, ForwardDiff, Random, SparseDiffTools
julia> function gnn_test()
Random.seed!(1234)
g = GNNGraph(erdos_renyi(10, 30), ndata=rand(Float32, 3, 10))
m = GATConv(3 => 2, tanh)
ps, re = Flux.destructure(m) # primal vector and restructure function
ts = rand(Float32, size(ps)) # tangent vector
loss(_ps) = sum(re(_ps)(g, g.ndata.x))
numback_hesvec(loss, ps, ts) |> println # works
numback_hesvec(loss, ps, ts) |> println # works
numauto_hesvec(loss, ps, ts) |> println # works
autoback_hesvec(loss, ps, ts) |> println # fails for GCNConv, works for GATConv
end
gnn_test (generic function with 1 method)
julia> gnn_test()
Float32[4.4092455, -0.75034475, 2.9557762, -0.39457783, 2.5744586, 0.34606418, 8.512531, 0.21992864, -0.2590933, -0.5618502, 0.4025826, -1.1103446]
Float32[4.4092455, -0.75034475, 2.9557762, -0.39457783, 2.5744586, 0.34606418, 8.512531, 0.21992864, -0.2590933, -0.5618502, 0.4025826, -1.1103446]
Float32[4.4092455, -0.751315, 2.9551294, -0.39425442, 2.5746205, 0.34541732, 8.512531, 0.21992864, -0.25904277, -0.56184894, 0.40250173, -1.1103705]
Float32[4.409217, -0.7511167, 2.9556952, -0.3946823, 2.5747547, 0.34566167, 8.511975, 0.21927829, -0.25907597, -0.5618338, 0.40243906, -1.1103226]
but doesn't like operations involving sparse matrices (i.e. what happens for GCNConv`).
The problem seems a very generic one not strictly related to GNN.jl. Here is a MWE:
julia> using SparseArrays, SparseDiffTools
julia> x, t = rand(5), rand(5);
julia> A = sprand(5, 5, 0.5);
julia> loss(x) = sum(tanh.(A * x));
julia> numback_hesvec(loss, x, t) # works
5-element Vector{Float64}:
-0.349703846209146
-1.210662833747414
-1.4030571895355597
-0.47786341057923254
-0.9171474544184983
julia> autoback_hesvec(loss, x, t)
ERROR: MethodError: no method matching Float64(::ForwardDiff.Dual{ForwardDiff.Tag{DataType, Float64}, Float64, 1})
Closest candidates are:
(::Type{T})(::Real, ::RoundingMode) where T<:AbstractFloat at ~/julia/julia-1.7.1/share/julia/base/rounding.jl:200
(::Type{T})(::T) where T<:Number at ~/julia/julia-1.7.1/share/julia/base/boot.jl:770
(::Type{T})(::AbstractChar) where T<:Union{AbstractChar, Number} at ~/julia/julia-1.7.1/share/julia/base/char.jl:50
...
Stacktrace:
[1] convert(#unused#::Type{Float64}, x::ForwardDiff.Dual{ForwardDiff.Tag{DataType, Float64}, Float64, 1})
@ Base ./number.jl:7
[2] setindex!(A::Vector{Float64}, x::ForwardDiff.Dual{ForwardDiff.Tag{DataType, Float64}, Float64, 1}, i1::Int64)
@ Base ./array.jl:903
[3] (::ChainRulesCore.ProjectTo{SparseMatrixCSC, NamedTuple{(:element, :axes, :rowval, :nzranges, :colptr), Tuple{ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}, Vector{Int64}, Vector{UnitRange{Int64}}, Vector{Int64}}}})(dx::Matrix{ForwardDiff.Dual{ForwardDiff.Tag{DataType, Float64}, Float64, 1}})
@ ChainRulesCore ~/.julia/packages/ChainRulesCore/uxrij/src/projection.jl:580
[4] #1334
@ ~/.julia/packages/ChainRules/GRzER/src/rulesets/Base/arraymath.jl:36 [inlined]
[5] unthunk
@ ~/.julia/packages/ChainRulesCore/uxrij/src/tangent_types/thunks.jl:197 [inlined]
[6] wrap_chainrules_output
@ ~/.julia/packages/Zygote/FPUm3/src/compiler/chainrules.jl:104 [inlined]
[7] map
@ ./tuple.jl:223 [inlined]
[8] wrap_chainrules_output
@ ~/.julia/packages/Zygote/FPUm3/src/compiler/chainrules.jl:105 [inlined]
[9] ZBack
@ ~/.julia/packages/Zygote/FPUm3/src/compiler/chainrules.jl:204 [inlined]
[10] Pullback
@ ./REPL[40]:1 [inlined]
[11] (::typeof(∂(loss)))(Δ::ForwardDiff.Dual{ForwardDiff.Tag{DataType, Float64}, Float64, 1})
@ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
[12] (::Zygote.var"#57#58"{typeof(∂(loss))})(Δ::ForwardDiff.Dual{ForwardDiff.Tag{DataType, Float64}, Float64, 1})
@ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface.jl:41
[13] gradient(f::Function, args::Vector{ForwardDiff.Dual{ForwardDiff.Tag{DataType, Float64}, Float64, 1}})
@ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface.jl:76
[14] (::SparseDiffTools.var"#78#79"{typeof(loss)})(x::Vector{ForwardDiff.Dual{ForwardDiff.Tag{DataType, Float64}, Float64, 1}})
@ SparseDiffTools ~/.julia/packages/SparseDiffTools/b2cgD/src/differentiation/jaches_products_zygote.jl:39
[15] autoback_hesvec(f::Function, x::Vector{Float64}, v::Vector{Float64})
@ SparseDiffTools ~/.julia/packages/SparseDiffTools/b2cgD/src/differentiation/jaches_products_zygote.jl:41
[16] top-level scope
@ REPL[47]:1
[17] top-level scope
@ ~/.julia/packages/CUDA/bki2w/src/initialization.jl:52
Hi think this issue should be reported to SparseDiffTools.jl and probably fixed in ChainRules.jl
Hello! Nice work on the library; it is very usable. I'm trying to calculate the hessian-vector product of a loss function involving
GNNGraph
datapoints and aGNNChain
model. I've been using the SparseDiffTools.jl functionautoback_hesvec
for this, which implements ForwardDiff.jl over Zygote.jl for the hessian-vector calculation. However, this function is failing in the GraphNeuralNetworks.jl setting. The other hessian-vector functions in SparseDiffTools.jl do work, and an analogously-constructed calculation using only Flux works.The full error message: