From 331b21bf9742f885b2037d1b6c7cdc19d757b733 Mon Sep 17 00:00:00 2001 From: Grant Bruer Date: Fri, 12 Sep 2025 15:20:42 -0400 Subject: [PATCH] Added hyperbolic trig functions for activations --- .../conditional_layer_glow.jl | 6 +- src/layers/invertible_layer_basic.jl | 26 +- src/layers/invertible_layer_glow.jl | 14 +- src/layers/layer_residual_block.jl | 21 +- src/utils/activation_functions.jl | 90 +++++- test/test_layers/test_residual_block.jl | 14 +- test/test_networks/test_glow.jl | 294 +++++++++--------- ...est_multiscale_conditional_hint_network.jl | 18 +- 8 files changed, 293 insertions(+), 190 deletions(-) diff --git a/src/conditional_layers/conditional_layer_glow.jl b/src/conditional_layers/conditional_layer_glow.jl index 5fb6c5cd..86abe9e8 100644 --- a/src/conditional_layers/conditional_layer_glow.jl +++ b/src/conditional_layers/conditional_layer_glow.jl @@ -136,14 +136,14 @@ function inverse(Y::AbstractArray{T, N}, C::AbstractArray{T, N}, L::ConditionalL X_ = tensor_cat(X1, X2) X = L.C.inverse(X_) - save == true ? (return X, X1, X2, Sm) : (return X) + save == true ? (return X, X1, X2, logS, Sm) : (return X) end # Backward pass: Input (ΔY, Y), Output (ΔX, X) function backward(ΔY::AbstractArray{T, N}, Y::AbstractArray{T, N}, C::AbstractArray{T, N}, L::ConditionalLayerGlow;) where {T,N} # Recompute forward state - X, X1, X2, S = inverse(Y, C, L; save=true) + X, X1, X2, logS, S = inverse(Y, C, L; save=true) # Backpropagate residual ΔY1, ΔY2 = tensor_split(ΔY) @@ -159,7 +159,7 @@ function backward(ΔY::AbstractArray{T, N}, Y::AbstractArray{T, N}, C::AbstractA end # Backpropagate RB - ΔX2_ΔC = L.RB.backward(tensor_cat(L.activation.backward(ΔS, S), ΔT), (tensor_cat(X2, C))) + ΔX2_ΔC = L.RB.backward(tensor_cat(apply_backward(L.activation, ΔS, logS, S), ΔT), (tensor_cat(X2, C))) ΔX2, ΔC = tensor_split(ΔX2_ΔC; split_index=size(ΔY2)[N-1]) ΔX2 += ΔY2 diff --git a/src/layers/invertible_layer_basic.jl b/src/layers/invertible_layer_basic.jl index 469e86af..4bc3743f 100644 --- a/src/layers/invertible_layer_basic.jl +++ b/src/layers/invertible_layer_basic.jl @@ -96,9 +96,9 @@ function forward(X1::AbstractArray{T, N}, X2::AbstractArray{T, N}, L::CouplingLa Y2 = S.*X2 + logS_T2 if logdet - save ? (return X1, Y2, coupling_logdet_forward(S), S) : (return X1, Y2, coupling_logdet_forward(S)) + save ? (return X1, Y2, coupling_logdet_forward(S), logS_T1, S) : (return X1, Y2, coupling_logdet_forward(S)) else - save ? (return X1, Y2, S) : (return X1, Y2) + save ? (return X1, Y2, logS_T1, S) : (return X1, Y2) end end @@ -112,9 +112,9 @@ function inverse(Y1::AbstractArray{T, N}, Y2::AbstractArray{T, N}, L::CouplingLa X2 = (Y2 - logS_T2) ./ (S .+ eps(T)) # add epsilon to avoid division by 0 if logdet - save == true ? (return Y1, X2, -coupling_logdet_forward(S), S) : (return Y1, X2, -coupling_logdet_forward(S)) + save == true ? (return Y1, X2, -coupling_logdet_forward(S), logS_T1, S) : (return Y1, X2, -coupling_logdet_forward(S)) else - save == true ? (return Y1, X2, S) : (return Y1, X2) + save == true ? (return Y1, X2, logS_T1, S) : (return Y1, X2) end end @@ -122,7 +122,7 @@ end function backward(ΔY1::AbstractArray{T, N}, ΔY2::AbstractArray{T, N}, Y1::AbstractArray{T, N}, Y2::AbstractArray{T, N}, L::CouplingLayerBasic; set_grad::Bool=true) where {T, N} # Recompute forward state - X1, X2, S = inverse(Y1, Y2, L; save=true, logdet=false) + X1, X2, logS_T1, S = inverse(Y1, Y2, L; save=true, logdet=false) # Backpropagate residual ΔT = copy(ΔY2) @@ -132,11 +132,11 @@ function backward(ΔY1::AbstractArray{T, N}, ΔY2::AbstractArray{T, N}, Y1::Abst end ΔX2 = ΔY2 .* S if set_grad - ΔX1 = L.RB.backward(tensor_cat(L.activation.backward(ΔS, S), ΔT), X1) + ΔY1 + ΔX1 = L.RB.backward(tensor_cat(apply_backward(L.activation, ΔS, logS_T1, S), ΔT), X1) + ΔY1 else - ΔX1, Δθ = L.RB.backward(tensor_cat(L.activation.backward(ΔS, S), ΔT), X1; set_grad=set_grad) + ΔX1, Δθ = L.RB.backward(tensor_cat(apply_backward(L.activation, ΔS, logS_T1, S), ΔT), X1; set_grad=set_grad) if L.logdet - _, ∇logdet = L.RB.backward(tensor_cat(L.activation.backward(coupling_logdet_backward(S), S), 0 .*ΔT), X1; set_grad=set_grad) + _, ∇logdet = L.RB.backward(tensor_cat(apply_backward(L.activation, coupling_logdet_backward(S), logS_T1, S), 0 .*ΔT), X1; set_grad=set_grad) end ΔX1 += ΔY1 end @@ -152,7 +152,7 @@ end function backward_inv(ΔX1::AbstractArray{T, N}, ΔX2::AbstractArray{T, N}, X1::AbstractArray{T, N}, X2::AbstractArray{T, N}, L::CouplingLayerBasic; set_grad::Bool=true) where {T, N} # Recompute inverse state - Y1, Y2, S = forward(X1, X2, L; save=true, logdet=false) + Y1, Y2, logS_T1, S = forward(X1, X2, L; save=true, logdet=false) # Backpropagate residual ΔT = -ΔX2 ./ S @@ -161,9 +161,9 @@ function backward_inv(ΔX1::AbstractArray{T, N}, ΔX2::AbstractArray{T, N}, X1:: set_grad ? (ΔS += coupling_logdet_backward(S)) : (∇logdet = -coupling_logdet_backward(S)) end if set_grad - ΔY1 = L.RB.backward(tensor_cat(L.activation.backward(ΔS, S), ΔT), Y1) + ΔX1 + ΔY1 = L.RB.backward(tensor_cat(apply_backward(L.activation, ΔS, logS_T1, S), ΔT), Y1) + ΔX1 else - ΔY1, Δθ = L.RB.backward(tensor_cat(L.activation.backward(ΔS, S), ΔT), Y1; set_grad=set_grad) + ΔY1, Δθ = L.RB.backward(tensor_cat(apply_backward(L.activation, ΔS, logS_T1, S), ΔT), Y1; set_grad=set_grad) ΔY1 += ΔX1 end ΔY2 = - ΔT @@ -187,14 +187,14 @@ function jacobian(ΔX1::AbstractArray{T, N}, ΔX2::AbstractArray{T, N}, Δθ::Ab logS_T1, logS_T2 = tensor_split(L.RB.forward(X1)) ΔlogS_T1, ΔlogS_T2 = tensor_split(jacobian(ΔX1, Δθ, X1, L.RB)[1]) S = L.activation.forward(logS_T1) - ΔS = L.activation.backward(ΔlogS_T1, S) + ΔS = apply_backward(L.activation, ΔlogS_T1, logS_T1, S) Y2 = S.*X2 + logS_T2 ΔY2 = ΔS.*X2 + S.*ΔX2 + ΔlogS_T2 if logdet # Gauss-Newton approximation of logdet terms JΔθ = tensor_split(L.RB.jacobian(zeros(Float32, size(ΔX1)), Δθ, X1)[1])[1] - GNΔθ = -L.RB.adjointJacobian(tensor_cat(L.activation.backward(JΔθ, S), zeros(Float32, size(S))), X1)[2] + GNΔθ = -L.RB.adjointJacobian(tensor_cat(apply_backward(L.activation, JΔθ, logS_T1, S), zeros(Float32, size(S))), X1)[2] save ? (return ΔX1, ΔY2, X1, Y2, coupling_logdet_forward(S), GNΔθ, S) : (return ΔX1, ΔY2, X1, Y2, coupling_logdet_forward(S), GNΔθ) else diff --git a/src/layers/invertible_layer_glow.jl b/src/layers/invertible_layer_glow.jl index 7a2c70fd..fc6405aa 100644 --- a/src/layers/invertible_layer_glow.jl +++ b/src/layers/invertible_layer_glow.jl @@ -129,14 +129,14 @@ function inverse(Y::AbstractArray{T, N}, L::CouplingLayerGlow; save=false) where X_ = tensor_cat(X1, X2) X = L.C.inverse(X_) - save == true ? (return X, X1, X2, Sm) : (return X) + save == true ? (return X, X1, X2, logSm, Sm) : (return X) end # Backward pass: Input (ΔY, Y), Output (ΔX, X) function backward(ΔY::AbstractArray{T, N}, Y::AbstractArray{T, N}, L::CouplingLayerGlow; set_grad::Bool=true) where {T,N} # Recompute forward state - X, X1, X2, S = inverse(Y, L; save=true) + X, X1, X2, logSm, S = inverse(Y, L; save=true) # Backpropagate residual ΔY1, ΔY2 = tensor_split(ΔY) @@ -148,10 +148,10 @@ function backward(ΔY::AbstractArray{T, N}, Y::AbstractArray{T, N}, L::CouplingL ΔX1 = ΔY1 .* S if set_grad - ΔX2 = L.RB.backward(tensor_cat(L.activation.backward(ΔS, S), ΔT), X2) + ΔY2 + ΔX2 = L.RB.backward(tensor_cat(apply_backward(L.activation, ΔS, logSm, S), ΔT), X2) + ΔY2 else - ΔX2, Δθrb = L.RB.backward(tensor_cat(L.activation.backward(ΔS, S), ΔT; ), X2; set_grad=set_grad) - _, ∇logdet = L.RB.backward(tensor_cat(L.activation.backward(ΔS, S), 0f0.*ΔT;), X2; set_grad=set_grad) + ΔX2, Δθrb = L.RB.backward(tensor_cat(apply_backward(L.activation, ΔS, logSm, S), ΔT; ), X2; set_grad=set_grad) + _, ∇logdet = L.RB.backward(tensor_cat(apply_backward(L.activation, ΔS, logSm, S), 0f0.*ΔT;), X2; set_grad=set_grad) ΔX2 += ΔY2 end ΔX_ = tensor_cat(ΔX1, ΔX2) @@ -187,7 +187,7 @@ function jacobian(ΔX::AbstractArray{T, N}, Δθ::Array{Parameter, 1}, X, L::Cou ΔlogS, ΔlogT = tensor_split(ΔlogS_T) logS, logT = tensor_split(logS_T) Sm = L.activation.forward(logS) - ΔS = L.activation.backward(ΔlogS, nothing;x=logS) + ΔS = apply_backward(L.activation, ΔlogS, logS, Sm) Tm = logT ΔT = ΔlogT Y1 = Sm.*X1 + Tm @@ -197,7 +197,7 @@ function jacobian(ΔX::AbstractArray{T, N}, Δθ::Array{Parameter, 1}, X, L::Cou # Gauss-Newton approximation of logdet terms JΔθ,_ = tensor_split(L.RB.jacobian(cuzeros(ΔX2, size(ΔX2)), Δθ[4:end], X2)[1])#[:, :, 1:k, :] - GNΔθ = cat(0f0*Δθ[1:3], -L.RB.adjointJacobian(tensor_cat(L.activation.backward(JΔθ, Sm), zeros(Float32, size(Sm))), X2)[2]; dims=1) + GNΔθ = cat(0f0*Δθ[1:3], -L.RB.adjointJacobian(tensor_cat(apply_backward(L.activation, JΔθ, logS, Sm), zeros(Float32, size(Sm))), X2)[2]; dims=1) L.logdet ? (return ΔY, Y, glow_logdet_forward(Sm), GNΔθ) : (return ΔY, Y) end diff --git a/src/layers/layer_residual_block.jl b/src/layers/layer_residual_block.jl index 60fb58f3..6cfccfd0 100644 --- a/src/layers/layer_residual_block.jl +++ b/src/layers/layer_residual_block.jl @@ -127,10 +127,13 @@ function forward(X1::AbstractArray{T, N}, RB::ResidualBlock; save=false) where { cdims3 = DCDims(X1, RB.W3.data; stride=RB.strides[1], padding=RB.pad[1]) Y3 = ∇conv_data(X3, RB.W3.data, cdims3) + # Return if only recomputing state - save && (return Y1, Y2, Y3) + X4 = RB.fan == true ? RB.activation.forward(Y3) : GaLU(Y3) + save && (return Y1, Y2, Y3, X2, X3, X4) + # Finish forward - RB.fan == true ? (return RB.activation.forward(Y3)) : (return GaLU(Y3)) + return X4 end # Backward @@ -140,25 +143,25 @@ function backward(ΔX4::AbstractArray{T, N}, X1::AbstractArray{T, N}, dims = collect(1:N-1); dims[end] +=1 # Recompute forward states from input X - Y1, Y2, Y3 = forward(X1, RB; save=true) + Y1, Y2, Y3, X2, X3, X4 = forward(X1, RB; save=true) # Cdims cdims2 = DenseConvDims(Y2, RB.W2.data; stride=RB.strides[2], padding=RB.pad[2]) cdims3 = DCDims(X1, RB.W3.data; stride=RB.strides[1], padding=RB.pad[1]) # Backpropagate residual ΔX4 and compute gradients - RB.fan == true ? (ΔY3 = RB.activation.backward(ΔX4, Y3)) : (ΔY3 = GaLUgrad(ΔX4, Y3)) + RB.fan == true ? (ΔY3 = apply_backward(RB.activation, ΔX4, Y3, X4)) : (ΔY3 = GaLUgrad(ΔX4, Y3)) ΔX3 = conv(ΔY3, RB.W3.data, cdims3) ΔW3 = ∇conv_filter(ΔY3, RB.activation.forward(Y2), cdims3) - ΔY2 = RB.activation.backward(ΔX3, Y2) + ΔY2 = apply_backward(RB.activation, ΔX3, Y2, X3) ΔX2 = ∇conv_data(ΔY2, RB.W2.data, cdims2) + ΔY2 ΔW2 = ∇conv_filter(RB.activation.forward(Y1), ΔY2, cdims2) Δb2 = sum(ΔY2, dims=dims)[inds...] cdims1 = DenseConvDims(X1, RB.W1.data; stride=RB.strides[1], padding=RB.pad[1]) - ΔY1 = RB.activation.backward(ΔX2, Y1) + ΔY1 = apply_backward(RB.activation, ΔX2, Y1, X2) ΔX1 = ∇conv_data(ΔY1, RB.W1.data, cdims1) ΔW1 = ∇conv_filter(X1, ΔY1, cdims1) Δb1 = sum(ΔY1, dims=dims)[inds...] @@ -187,21 +190,21 @@ function jacobian(ΔX1::AbstractArray{T, N}, Δθ::Array{Parameter, 1}, Y1 = conv(X1, RB.W1.data, cdims1) .+ reshape(RB.b1.data, inds...) ΔY1 = conv(ΔX1, RB.W1.data, cdims1) + conv(X1, Δθ[1].data, cdims1) .+ reshape(Δθ[4].data, inds...) X2 = RB.activation.forward(Y1) - ΔX2 = RB.activation.backward(ΔY1, Y1) + ΔX2 = apply_backward(RB.activation, ΔY1, Y1, X2) cdims2 = DenseConvDims(X2, RB.W2.data; stride=RB.strides[2], padding=RB.pad[2]) Y2 = X2 + conv(X2, RB.W2.data, cdims2) .+ reshape(RB.b2.data, inds...) ΔY2 = ΔX2 + conv(ΔX2, RB.W2.data, cdims2) + conv(X2, Δθ[2].data, cdims2) .+ reshape(Δθ[5].data, inds...) X3 = RB.activation.forward(Y2) - ΔX3 = RB.activation.backward(ΔY2, Y2) + ΔX3 = apply_backward(RB.activation, ΔY2, Y2, X3) cdims3 = DCDims(X1, RB.W3.data; nc=2*size(X1, N-1), stride=RB.strides[1], padding=RB.pad[1]) Y3 = ∇conv_data(X3, RB.W3.data, cdims3) ΔY3 = ∇conv_data(ΔX3, RB.W3.data, cdims3) + ∇conv_data(X3, Δθ[3].data, cdims3) if RB.fan == true X4 = RB.activation.forward(Y3) - ΔX4 = RB.activation.backward(ΔY3, Y3) + ΔX4 = apply_backward(RB.activation, ΔY3, Y3, X4) else ΔX4, X4 = GaLUjacobian(ΔY3, Y3) end diff --git a/src/utils/activation_functions.jl b/src/utils/activation_functions.jl index 8695f23c..afb6b5d7 100644 --- a/src/utils/activation_functions.jl +++ b/src/utils/activation_functions.jl @@ -8,6 +8,8 @@ export Sigmoid, SigmoidInv, SigmoidGrad export GaLU, GaLUgrad export ExpClamp, ExpClampInv, ExpClampGrad export ReLUlayer, LeakyReLUlayer, SigmoidLayer, Sigmoid2Layer, GaLUlayer, ExpClampLayer +export IdentityActivation, SoftplusLayer, TanhLayer, CoshLayer, SinhLayer +export apply_backward ############################################################################### @@ -19,6 +21,25 @@ struct ActivationFunction backward::Function end +""" +Helper function for cases where the caller does not know if the activation function is invertible. +""" +function apply_backward(activation::ActivationFunction, Δy::AbstractArray{T, N}, x::AbstractArray{T, N}, y::AbstractArray{T, N}) where {T, N} + apply_backward(activation.backward, activation.inverse, Δy, x, y) +end + +function apply_backward(backward::Function, inverse::Nothing, Δy::AbstractArray{T, N}, x::AbstractArray{T, N}, y::AbstractArray{T, N}) where {T, N} + backward(Δy, x) +end + +function apply_backward(backward::Function, inverse::Function, Δy::AbstractArray{T, N}, x::AbstractArray{T, N}, y::AbstractArray{T, N}) where {T, N} + backward(Δy, y) +end + + +IdentityActivation() = ActivationFunction(identity, identity, IdentityGrad) +IdentityGrad(Δy::AbstractArray{T, N}, x::AbstractArray{T, N}) where {T, N} = Δy + function ReLUlayer() return ActivationFunction(ReLU, nothing, ReLUgrad) end @@ -46,7 +67,7 @@ function GaLUlayer() end function ExpClampLayer() - return ActivationFunction(x -> ExpClamp(x), y -> ExpClampInv(y/2f0), (Δy, y) -> ExpClampGrad(Δy*2f0, y/2f0)) + return ActivationFunction(x -> 2 * ExpClamp(x), y -> ExpClampInv(y/2), (Δy, y) -> ExpClampGrad(Δy*2, y/2)) end @@ -308,4 +329,69 @@ function ExpClampGrad(Δy::AbstractArray{T, N}, y::AbstractArray{T, N}; x=nothin return clamp * T(0.636) * Δy .* y ./ (1 .+ x.^2) end -ExpClampGrad(Δy::AbstractArray{T, N}, ::Nothing; x=nothing, clamp=T(2)) where {T, N} = clamp * T(0.636) * Δy .* y ./ (1 .+ x.^2) \ No newline at end of file +ExpClampGrad(Δy::AbstractArray{T, N}, ::Nothing; x=nothing, clamp=T(2)) where {T, N} = clamp * T(0.636) * Δy .* y ./ (1 .+ x.^2) + + + +SoftplusLayer() = ActivationFunction(Softplus, SoftplusInv, SoftplusGrad) + +function Softplus(x::AbstractArray{T, N}) where {T, N} + return log.(1 .+ exp.(x)) +end + +function SoftplusInv(y::AbstractArray{T, N}) where {T, N} + if any(y .≈ 0) + throw(InputError("Input contains zeros.")) + else + return log.(exp.(y) .- 1) + end +end + +function SoftplusGrad(Δy::AbstractArray{T, N}, y::AbstractArray{T, N}) where {T, N} + return (exp.(y) .- 1) ./ exp.(y) .* Δy +end + + +TanhLayer() = ActivationFunction(Tanh, TanhInv, TanhGrad) + +function Tanh(x::AbstractArray{T, N}) where {T, N} + return tanh.(x) +end + +function TanhInv(y::AbstractArray{T, N}) where {T, N} + if any(abs.(y) .> 1 - 1f-6) + throw(InputError("Input outside tanh range.")) + else + return atanh.(y) + end +end + +function TanhGrad(Δy::AbstractArray{T, N}, y::AbstractArray{T, N}) where {T, N} + return (1 .- y .^ 2) .* Δy +end + + +CoshLayer() = ActivationFunction(Cosh, nothing, CoshGrad) + +function Cosh(x::AbstractArray{T, N}) where {T, N} + return cosh.(x) +end + +function CoshGrad(Δy::AbstractArray{T, N}, x::AbstractArray{T, N}) where {T, N} + return sinh.(x) .* Δy +end + + +SinhLayer() = ActivationFunction(Sinh, SinhInv, SinhGrad) + +function Sinh(x::AbstractArray{T, N}) where {T, N} + return sinh.(x) +end + +function SinhInv(y::AbstractArray{T, N}) where {T, N} + return asinh.(y) +end + +function SinhGrad(Δy::AbstractArray{T, N}, y::AbstractArray{T, N}) where {T, N} + return cosh.(asinh.(y)) .* Δy +end diff --git a/test/test_layers/test_residual_block.jl b/test/test_layers/test_residual_block.jl index b8299fbb..4169bac5 100644 --- a/test/test_layers/test_residual_block.jl +++ b/test/test_layers/test_residual_block.jl @@ -15,8 +15,18 @@ k1 = 3 k2 = 3 - -for activation in [ReLUlayer(), LeakyReLUlayer()] +for activation in [ + ReLUlayer(), + LeakyReLUlayer(), + SigmoidLayer(), + Sigmoid2Layer(), + ExpClampLayer(), + IdentityActivation(), + SoftplusLayer(), + TanhLayer(), + CoshLayer(), + SinhLayer(), +] println("Testing activation $(activation)") # Input X = glorot_uniform(nx, ny, n_in, batchsize); diff --git a/test/test_networks/test_glow.jl b/test/test_networks/test_glow.jl index 0a0e1af1..41539c3a 100644 --- a/test/test_networks/test_glow.jl +++ b/test/test_networks/test_glow.jl @@ -29,175 +29,177 @@ K = 2 for logdet_bool = [true,false] #logdet is not tested in Jacobian yet. for split_scales = [true,false] for N in [(nx, ny), (nx, ny, nz)] - println("Testing Glow with dimensions=$(N) logdet=$(logdet_bool) and split_scales=$(split_scales)") - - # Network and input - G = NetworkGlow(n_in, n_hidden, L, K;logdet=logdet_bool, split_scales=split_scales, ndims=length(N))|> device - X = rand(m,Float32, N..., n_in, batchsize)|> device - - # Invertibility - if logdet_bool - Y, _ = G.forward(X) - else - Y = G.forward(X) - end - X_ = G.inverse(Y) + for activation in [SigmoidLayer(), CoshLayer()] + println("Testing Glow with dimensions=$(N) logdet=$(logdet_bool) and split_scales=$(split_scales)") - @test isapprox(norm(X - X_)/norm(X), 0f0; atol=1f-5) + # Network and input + G = NetworkGlow(n_in, n_hidden, L, K;logdet=logdet_bool, split_scales=split_scales, ndims=length(N), activation=activation)|> device + X = rand(m,Float32, N..., n_in, batchsize)|> device - ################################################################################################### - # Test gradients are set and cleared - G.backward(Y, Y) + # Invertibility + if logdet_bool + Y, _ = G.forward(X) + else + Y = G.forward(X) + end + X_ = G.inverse(Y) - P = get_params(G) - gsum = 0 - for p in P - ~isnothing(p.grad) && (gsum += 1) - end - - param_factor = 10 - @test isequal(gsum, L*K*param_factor) - - clear_grad!(G) - gsum = 0 - for p in P - ~isnothing(p.grad) && (gsum += 1) - end - @test isequal(gsum, 0) + @test isapprox(norm(X - X_)/norm(X), 0f0; atol=1f-5) + ################################################################################################### + # Test gradients are set and cleared + G.backward(Y, Y) - ################################################################################################### - # Gradient test + P = get_params(G) + gsum = 0 + for p in P + ~isnothing(p.grad) && (gsum += 1) + end - - function loss(G, X) - if G.logdet - Y, logdet = G.forward(X) - f = -log_likelihood(Y) - logdet - else - Y = G.forward(X) - f = -log_likelihood(Y) - end - ΔY = -∇log_likelihood(Y) - ΔX, X_ = G.backward(ΔY, Y) - return f, ΔX, G.CL[1,1].RB.W1.grad, G.CL[1,1].C.v1.grad - end + param_factor = 10 + @test isequal(gsum, L*K*param_factor) - # Gradient test w.r.t. input - G = NetworkGlow(n_in, n_hidden, L, K;logdet=logdet_bool, split_scales=split_scales, ndims=length(N))|> device - X = rand(m,Float32, N..., n_in, batchsize)|> device - X0 = rand(m,Float32, N..., n_in, batchsize)|> device - dX = X - X0 - - f0, ΔX = loss(G, X0)[1:2] - h = 0.1f0 - maxiter = 4 - err1 = zeros(Float32, maxiter) - err2 = zeros(Float32, maxiter) - - print("\nGradient test glow: input\n") - for j=1:maxiter - f = loss(G, X0 + h*dX,)[1] - err1[j] = abs(f - f0) - err2[j] = abs(f - f0 - h*dot(dX, ΔX)) - print(err1[j], "; ", err2[j], "\n") - h = h/2f0 - end + clear_grad!(G) + gsum = 0 + for p in P + ~isnothing(p.grad) && (gsum += 1) + end + @test isequal(gsum, 0) - @test isapprox(err1[end] / (err1[1]/2^(maxiter-1)), 1f0; atol=1f1) - @test isapprox(err2[end] / (err2[1]/4^(maxiter-1)), 1f0; atol=1f1) - num_attempts = 3 - results_1 = [] - results_2 = [] - for i in 1:num_attempts - Random.seed!(i) + ################################################################################################### + # Gradient test - # Gradient test w.r.t. parameters - X = rand(m,Float32, N..., n_in, batchsize)|> device - G = NetworkGlow(n_in, n_hidden, L, K; logdet=logdet_bool, split_scales=split_scales, ndims=length(N))|> device - G0 = NetworkGlow(n_in, n_hidden, L, K;logdet=logdet_bool, split_scales=split_scales, ndims=length(N))|> device - Gini = deepcopy(G0) - # Test one parameter from residual block and 1x1 conv - dW = G.CL[1,1].RB.W1.data - G0.CL[1,1].RB.W1.data - dv = G.CL[1,1].C.v1.data - G0.CL[1,1].C.v1.data + function loss(G, X) + if G.logdet + Y, logdet = G.forward(X) + f = -log_likelihood(Y) - logdet + else + Y = G.forward(X) + f = -log_likelihood(Y) + end + ΔY = -∇log_likelihood(Y) + ΔX, X_ = G.backward(ΔY, Y) + return f, ΔX, G.CL[1,1].RB.W1.grad, G.CL[1,1].C.v1.grad + end + + # Gradient test w.r.t. input + G = NetworkGlow(n_in, n_hidden, L, K;logdet=logdet_bool, split_scales=split_scales, ndims=length(N), activation=activation)|> device + X = rand(m,Float32, N..., n_in, batchsize)|> device + X0 = rand(m,Float32, N..., n_in, batchsize)|> device + dX = X - X0 - f0, ΔX, ΔW, Δv = loss(G0, X) + f0, ΔX = loss(G, X0)[1:2] h = 0.1f0 maxiter = 4 - err3 = zeros(Float32, maxiter) - err4 = zeros(Float32, maxiter) + err1 = zeros(Float32, maxiter) + err2 = zeros(Float32, maxiter) print("\nGradient test glow: input\n") for j=1:maxiter - G0.CL[1,1].RB.W1.data = Gini.CL[1,1].RB.W1.data + h*dW - G0.CL[1,1].C.v1.data = Gini.CL[1,1].C.v1.data + h*dv - - f = loss(G0, X)[1] - err3[j] = abs(f - f0) - err4[j] = abs(f - f0 - h*dot(dW, ΔW) - h*dot(dv, Δv)) - print(err3[j], "; ", err4[j], "\n") - h = h*hfactor + f = loss(G, X0 + h*dX,)[1] + err1[j] = abs(f - f0) + err2[j] = abs(f - f0 - h*dot(dX, ΔX)) + print(err1[j], "; ", err2[j], "\n") + h = h/2f0 end - factor1 = err3[1:end-1]./err3[2:end] - factor2 = err4[1:end-1]./err4[2:end] + @test isapprox(err1[end] / (err1[1]/2^(maxiter-1)), 1f0; atol=1f1) + @test isapprox(err2[end] / (err2[1]/4^(maxiter-1)), 1f0; atol=1f1) + + num_attempts = 3 + results_1 = [] + results_2 = [] + for i in 1:num_attempts + Random.seed!(i) + + # Gradient test w.r.t. parameters + X = rand(m,Float32, N..., n_in, batchsize)|> device + G = NetworkGlow(n_in, n_hidden, L, K; logdet=logdet_bool, split_scales=split_scales, ndims=length(N), activation=activation)|> device + G0 = NetworkGlow(n_in, n_hidden, L, K;logdet=logdet_bool, split_scales=split_scales, ndims=length(N), activation=activation)|> device + Gini = deepcopy(G0) + + # Test one parameter from residual block and 1x1 conv + dW = G.CL[1,1].RB.W1.data - G0.CL[1,1].RB.W1.data + dv = G.CL[1,1].C.v1.data - G0.CL[1,1].C.v1.data + + f0, ΔX, ΔW, Δv = loss(G0, X) + h = 0.1f0 + maxiter = 4 + err3 = zeros(Float32, maxiter) + err4 = zeros(Float32, maxiter) + + print("\nGradient test glow: input\n") + for j=1:maxiter + G0.CL[1,1].RB.W1.data = Gini.CL[1,1].RB.W1.data + h*dW + G0.CL[1,1].C.v1.data = Gini.CL[1,1].C.v1.data + h*dv + + f = loss(G0, X)[1] + err3[j] = abs(f - f0) + err4[j] = abs(f - f0 - h*dot(dW, ΔW) - h*dot(dv, Δv)) + print(err3[j], "; ", err4[j], "\n") + h = h*hfactor + end + + factor1 = err3[1:end-1]./err3[2:end] + factor2 = err4[1:end-1]./err4[2:end] + + append!(results_1,isapprox(mean(factor1), expected_f1; atol=1f0)) + append!(results_2,isapprox(mean(factor2), expected_f2; atol=1f0)) + end + @test true in results_1 + @test true in results_2 - append!(results_1,isapprox(mean(factor1), expected_f1; atol=1f0)) - append!(results_2,isapprox(mean(factor2), expected_f2; atol=1f0)) - end - @test true in results_1 - @test true in results_2 - - ################################################################################################### - # Jacobian-related tests - - # Gradient test - - # Initialization - G = NetworkGlow(n_in, n_hidden, L, K; ndims=length(N))|> device; G.forward(randn(Float32, N..., n_in, batchsize)|> device) - θ = deepcopy(get_params(G)) - G0 = NetworkGlow(n_in, n_hidden, L, K; ndims=length(N))|> device; G0.forward(randn(Float32, N..., n_in, batchsize)|> device) - θ0 = deepcopy(get_params(G0)) - X = randn(Float32, N..., n_in, batchsize)|> device - - # Perturbation (normalized) - dθ = θ-θ0 - dθ .*= norm.(θ0)./(norm.(dθ).+1f-10) - dX = randn(Float32, N..., n_in, batchsize)|> device; dX *= norm(X)/norm(dX) - - - # Jacobian eval - dY, Y, _, _ = G.jacobian(dX|> device, dθ|> device, X) - - # Test - print("\nJacobian test\n") - h = 0.1f0 - maxiter = 5 - err5 = zeros(Float32, maxiter) - err6 = zeros(Float32, maxiter) - for j=1:maxiter - set_params!(G, θ+h*dθ) - Y_loc, _ = G.forward(X+h*dX) - err5[j] = norm(Y_loc - Y) - err6[j] = norm(Y_loc - Y - h*dY) - print(err5[j], "; ", err6[j], "\n") - h = h/2f0 - end + ################################################################################################### + # Jacobian-related tests + + # Gradient test + + # Initialization + G = NetworkGlow(n_in, n_hidden, L, K; ndims=length(N), activation=activation)|> device; G.forward(randn(Float32, N..., n_in, batchsize)|> device) + θ = deepcopy(get_params(G)) + G0 = NetworkGlow(n_in, n_hidden, L, K; ndims=length(N), activation=activation)|> device; G0.forward(randn(Float32, N..., n_in, batchsize)|> device) + θ0 = deepcopy(get_params(G0)) + X = randn(Float32, N..., n_in, batchsize)|> device + + # Perturbation (normalized) + dθ = θ-θ0 + dθ .*= norm.(θ0)./(norm.(dθ).+1f-10) + dX = randn(Float32, N..., n_in, batchsize)|> device; dX *= norm(X)/norm(dX) - @test isapprox(err5[end] / (err5[1]/2^(maxiter-1)), 1f0; atol=1f1) - @test isapprox(err6[end] / (err6[1]/4^(maxiter-1)), 1f0; atol=1f1) - # Adjoint test + # Jacobian eval + dY, Y, _, _ = G.jacobian(dX|> device, dθ|> device, X) - set_params!(G, θ) - dY, Y, _, _ = G.jacobian(dX, dθ, X) - dY_ = randn(Float32, size(dY))|> device - dX_, dθ_, _, _ = G.adjointJacobian(dY_, Y) - a = dot(dY, dY_) - b = dot(dX, dX_) + dot(dθ, dθ_) - @test isapprox(a, b; rtol=1f-2) + # Test + print("\nJacobian test\n") + h = 0.1f0 + maxiter = 5 + err5 = zeros(Float32, maxiter) + err6 = zeros(Float32, maxiter) + for j=1:maxiter + set_params!(G, θ+h*dθ) + Y_loc, _ = G.forward(X+h*dX) + err5[j] = norm(Y_loc - Y) + err6[j] = norm(Y_loc - Y - h*dY) + print(err5[j], "; ", err6[j], "\n") + h = h/2f0 + end + + @test isapprox(err5[end] / (err5[1]/2^(maxiter-1)), 1f0; atol=1f1) + @test isapprox(err6[end] / (err6[1]/4^(maxiter-1)), 1f0; atol=1f1) + + # Adjoint test + + set_params!(G, θ) + dY, Y, _, _ = G.jacobian(dX, dθ, X) + dY_ = randn(Float32, size(dY))|> device + dX_, dθ_, _, _ = G.adjointJacobian(dY_, Y) + a = dot(dY, dY_) + b = dot(dX, dX_) + dot(dθ, dθ_) + @test isapprox(a, b; rtol=1f-2) + end end end end diff --git a/test/test_networks/test_multiscale_conditional_hint_network.jl b/test/test_networks/test_multiscale_conditional_hint_network.jl index d94536aa..3ea6ebeb 100644 --- a/test/test_networks/test_multiscale_conditional_hint_network.jl +++ b/test/test_networks/test_multiscale_conditional_hint_network.jl @@ -14,9 +14,9 @@ batchsize = 2 L = 2 K = 2 -function inv_test(nx, ny, n_in, batchsize, logdet, squeeze_type, split_scales) - print("\nMultiscale Conditional HINT invertibility test with squeeze_type=$(squeeze_type), split_scales=$(split_scales), logdet=$(logdet)\n") - CH = NetworkMultiScaleConditionalHINT(n_in, n_hidden, L, K; squeezer = squeeze_type(), logdet=logdet, split_scales=split_scales) +function inv_test(nx, ny, n_in, batchsize, logdet, squeeze_type, split_scales, activation) + print("\nMultiscale Conditional HINT invertibility test with squeeze_type=$(squeeze_type), split_scales=$(split_scales), logdet=$(logdet), activation=$(activation)\n") + CH = NetworkMultiScaleConditionalHINT(n_in, n_hidden, L, K; squeezer = squeeze_type(), logdet=logdet, split_scales=split_scales, activation=activation) # Input image and data X = randn(Float32, nx, ny, n_in, batchsize) @@ -60,9 +60,9 @@ function loss(CH, X, Y) return f, ΔX, ΔY end -function grad_test_X(nx, ny, n_channel, batchsize, logdet, squeeze_type, split_scales) - print("\nMultiscale Conditional HINT gradient test with squeeze_type=$(squeeze_type), split_scales=$(split_scales), logdet=$(logdet)\n") - CH = NetworkMultiScaleConditionalHINT(n_in, n_hidden, L, K; squeezer = squeeze_type(), logdet=logdet, split_scales=split_scales) +function grad_test_X(nx, ny, n_channel, batchsize, logdet, squeeze_type, split_scales, activation) + print("\nMultiscale Conditional HINT gradient test with squeeze_type=$(squeeze_type), split_scales=$(split_scales), logdet=$(logdet), activation=$(activation)\n") + CH = NetworkMultiScaleConditionalHINT(n_in, n_hidden, L, K; squeezer = squeeze_type(), logdet=logdet, split_scales=split_scales, activation=activation) # Input image @@ -95,8 +95,10 @@ end for squeeze_i in [ShuffleLayer, WaveletLayer, HaarLayer] for split_scales in [true, false] for logdet in [false, true] - inv_test(nx, ny, n_in, batchsize, logdet, squeeze_i, split_scales) - grad_test_X(nx, ny, n_in, batchsize, logdet, squeeze_i, split_scales) + for activation in [SigmoidLayer(), CoshLayer()] + inv_test(nx, ny, n_in, batchsize, logdet, squeeze_i, split_scales, activation) + grad_test_X(nx, ny, n_in, batchsize, logdet, squeeze_i, split_scales, activation) + end end end end