Merge pull request #62 from ziyiyin97/master

Dispatch pullback rrule for all Invertible

Author: mloubout

src/utils/chainrules.jl

@@ -56,7 +56,7 @@ isa_newblock(state::InvertibleOperationsTape, X) = (state.counter_block == 0) ||
 """
 Error if mismatch between state and network
 """
-function check_coherence(state::InvertibleOperationsTape, net::Union{NeuralNetLayer,InvertibleNetwork})
+function check_coherence(state::InvertibleOperationsTape, net::Invertible)
     if state.counter_block != 0 && state.counter_layer != 0 && state.layer_blocks[state.counter_block][state.counter_layer] != net
         reset!(state)
         throw(ArgumentError("Current state does not correspond to current layer, resetting state..."))
@@ -66,7 +66,7 @@ end
 """
 Update state in the forward pass.
 """
-function forward_update!(state::InvertibleOperationsTape, X::AbstractArray{T,N}, Y::AbstractArray{T,N}, logdet::Union{Nothing,T}, net::Union{NeuralNetLayer,InvertibleNetwork}) where {T, N}
+function forward_update!(state::InvertibleOperationsTape, X::AbstractArray{T,N}, Y::AbstractArray{T,N}, logdet::Union{Nothing,T}, net::Invertible) where {T, N}
 
     if isa_newblock(state, X)
         push!(state.Y, Y)
@@ -104,7 +104,7 @@ end
 
 ## Chain rules for invertible networks
 # General pullback function
-function pullback(net::Union{NeuralNetLayer,InvertibleNetwork}, ΔY::AbstractArray{T,N};
+function pullback(net::Invertible, ΔY::AbstractArray{T,N};
                  state::InvertibleOperationsTape=GLOBAL_STATE_INVOPS) where {T, N}
 
     # Check state coherency
@@ -124,7 +124,7 @@ end
 
 
 # Reverse-mode AD rule
-function ChainRulesCore.rrule(net::Union{NeuralNetLayer,InvertibleNetwork}, X::AbstractArray{T, N};
+function ChainRulesCore.rrule(net::Invertible, X::AbstractArray{T, N};
                               state::InvertibleOperationsTape=GLOBAL_STATE_INVOPS) where {T, N}
 
     # Forward pass

src/utils/invertible_network_sequential.jl

@@ -6,7 +6,7 @@ export ComposedInvertibleNetwork, Composition
 import Base.length, Base.∘
 
 struct ComposedInvertibleNetwork <: InvertibleNetwork
-    layers::Array{T, 1} where {T <: Union{NeuralNetLayer, InvertibleNetwork}}
+    layers::Array{T, 1} where {T <: Invertible}
     logdet_array::Array{Bool, 1}
     logdet::Bool
     npars::Array{Int64, 1}
@@ -21,7 +21,7 @@ function Composition(layer...)
 
     # Initializing output
     depth = length(layer)
-    net_array = Array{Union{NeuralNetLayer, InvertibleNetwork}, 1}(undef, depth)
+    net_array = Array{Invertible, 1}(undef, depth)
     logdet_array = Array{Bool, 1}(undef, depth)
     logdet = false
     npars = Array{Int64, 1}(undef, depth)

src/utils/jacobian.jl

@@ -15,7 +15,7 @@ struct JacobianInvertibleNetwork{T} <: joAbstractLinearOperator{T, T}
     m::Int64
     fop::Function
     fop_T::Function
-    N::Union{NeuralNetLayer, InvertibleNetwork}
+    N::Invertible
     X::AbstractArray{T}
     Y::Union{Nothing, AbstractArray{T}}
 end

src/utils/neuralnet.jl

@@ -108,7 +108,6 @@ reset!(AI::Array{<:Invertible}) = for I ∈ AI reset!(I) end
  Resets the gradient of all the parameters in NL
 """
 clear_grad!(I::Invertible) = clear_grad!(get_params(I))
-clear_grad!(RL::Reversed) = clear_grad!(RL.I)
 
 # Get gradients
 """
@@ -124,12 +123,9 @@ get_grads(RL::Reversed)= get_grads(RL.I)
 get_grads(::Nothing) = []
 
 # Set parameters
-function set_params!(N::Union{NeuralNetLayer, InvertibleNetwork}, θnew::Array{Parameter, 1})
+function set_params!(N::Invertible, θnew::Array{Parameter, 1})
     set_params!(get_params(N), θnew)
 end
 
-# Set params for reversed layers/networks
-set_params!(RL::Reversed, θ::Array{Parameter, 1}) = set_params!(RL.I, θ)
-
 # Make invertible nets callable objects
 (N::Invertible)(X::AbstractArray{T,N} where {T, N}) = N.forward(X)

test/test_utils/test_chainrules.jl

@@ -60,4 +60,14 @@ N = Chain(N1, N2, N3, N4, N5, N6, N7, N8, N9, N10);
 
 g2 = gradient(X -> loss(X), X)
 
-@test g ≈ g2[1] rtol=1f-6
+@test g ≈ g2[1] rtol=1f-6
+
+## test Reverse network AD
+
+Nrev = reverse(N10)
+Xrev, ∂rev = rrule(Nrev, X)
+grev = ∂rev(Xrev-Y0)
+
+g2rev = gradient(X -> 0.5f0*norm(Nrev(X) - Y0)^2, X)
+
+@test grev[2] ≈ g2rev[1] rtol=1f-6