Move to ConditionalDists v0.2

Project.toml

@@ -1,12 +1,13 @@
 name = "GenerativeModels"
 uuid = "6ac2c632-c4cd-11e9-0501-33c4b9b2f9c9"
 authors = ["Niklas Heim <[email protected]>"]
-version = "0.1.0"
+version = "0.1.1"
 
 [deps]
 BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0"
 ConditionalDists = "c648c4dd-c1e0-49a6-84b9-144ae7fd2468"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DrWatson = "634d3b9d-ee7a-5ddf-bec9-22491ea816e1"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
@@ -19,8 +20,9 @@ ValueHistories = "98cad3c8-aec3-5f06-8e41-884608649ab7"
 
 [compat]
 BSON = "0.2"
-ConditionalDists = "0.1"
+ConditionalDists = "0.2"
 DiffEqBase = "6.15"
+Distributions = "0.23"
 DrWatson = "1.9"
 Flux = "0.10"
 ForwardDiff = "0.10"
@@ -30,13 +32,11 @@ Reexport = "0.2"
 ValueHistories = "0.5"
 julia = "1"
 
-
 [extras]
 CuArrays = "3a865a2d-5b23-5a0f-bc46-62713ec82fae"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"
-Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "Suppressor", "Logging", "Parameters", "CuArrays"]
+test = ["Test", "Logging", "Parameters", "CuArrays"]

src/GenerativeModels.jl

@@ -11,20 +11,19 @@ module GenerativeModels
     @reexport using IPMeasures
 
 
+    using Distributions: ContinuousMultivariateDistribution
+    using ConditionalDists: AbstractConditionalDistribution
+    const CMD = ContinuousMultivariateDistribution
+    const ACD = AbstractConditionalDistribution
+
     using Flux: @adjoint
     using DiffEqBase: ODEProblem, solve
     using OrdinaryDiffEq: Tsit5
-    using ConditionalDists: AbstractPDF, AbstractCPDF
 
     abstract type AbstractGM end
     abstract type AbstractVAE <: AbstractGM end
     abstract type AbstractGAN <: AbstractGM end
 
-    # functions that are overloaded by this module
-    import Base.length
-    import Random.rand
-    import Statistics.mean
-
     include(joinpath("utils", "flux_ode_decoder.jl"))
     include(joinpath("utils", "saveload.jl"))
     include(joinpath("utils", "utils.jl"))

src/models/gan.jl

@@ -2,7 +2,7 @@ export GAN
 export generator_loss, discriminator_loss
 
 """
-    GAN{P<:AbstractPDF,G<:AbstractCPDF,D<:AbstractCPDF}([zlength::Int,p::P], g::G, d::D)
+    GAN{P<:CMD,G<:ACD,D<:ACD}([zlength::Int,p::P], g::G, d::D)
 
 The Generative Adversarial Network.
 
@@ -28,7 +28,7 @@ GAN:
  discriminator = (CMeanGaussian{DiagVar}(mapping=Dense(4, 1, σ), σ2=1-element Array...)
 ```
 """
-struct GAN{P<:AbstractPDF,G<:AbstractCPDF,D<:AbstractCPDF} <: AbstractGAN
+struct GAN{P<:CMD,G<:ACD,D<:ACD} <: AbstractGAN
 	prior::P
 	generator::G
 	discriminator::D
@@ -38,7 +38,7 @@ Flux.@functor GAN
 
 GAN(p::P, g::G, d::D) where {P,G,D} = GAN{P,G,D}(p,g,d)
 
-function GAN(zlength::Int, g::AbstractCPDF, d::AbstractCPDF)
+function GAN(zlength::Int, g::ACD, d::ACD)
     T = eltype(first(params(g)))
     μ = NoGradArray(zeros(T, zlength))
     σ = NoGradArray(ones(T, zlength))

src/models/rodent.jl

@@ -93,7 +93,7 @@ function Rodent(slen::Int, tlen::Int, dt::T, encoder;
                 ode=Dense(slen,slen),
                 observe=sol->reshape(hcat(sol.u...), :),
                 olen=slen*tlen) where T
-    zlen = length(destructure(ode)) + slen
+    zlen = length(Flux.destructure(ode)[1]) + slen
 
     μpz = NoGradArray(zeros(T, zlen))
     λ2z = ones(T, zlen) / 20
@@ -126,7 +126,7 @@ function elbo(m::ConstSpecRodent, x::AbstractArray)
     sz = rand(m.encoder.spec, x)
     z  = cz .+ sz
 
-    llh = sum(loglikelihood(m.decoder, x, z))
+    llh = sum(logpdf(m.decoder, x, z))
     ckl = sum(kl_divergence(m.encoder.cnst, m.const_prior))
     skl = sum(kl_divergence(m.encoder.spec, m.spec_prior, sz))
 

src/models/vae.jl

@@ -2,7 +2,7 @@ export VAE
 export elbo, mmd_mean, mmd_rand
 
 """
-    VAE{P<:Gaussian,E<:AbstractCPDF,D<:AbstractCPDF}([zlength::Int,p::P] ,e::E ,d::D)
+    VAE{P<:Gaussian,E<:ACD,D<:ACD}([zlength::Int,p::P] ,e::E ,d::D)
 
 Variational Auto-Encoder.
 
@@ -36,7 +36,7 @@ julia> mean(vae.decoder, mean(vae.encoder, rand(5)))
   1.123661
 ```
 """
-struct VAE{P<:Gaussian,E<:AbstractCPDF,D<:AbstractCPDF} <: AbstractVAE
+struct VAE{P<:Gaussian,E<:ACD,D<:ACD} <: AbstractVAE
     prior::P
     encoder::E
     decoder::D
@@ -46,7 +46,7 @@ Flux.@functor VAE
 
 VAE(p::P, e::E, d::D) where {P,E,D} = VAE{P,E,D}(p,e,d)
 
-function VAE(zlength::Int, enc::AbstractCPDF, dec::AbstractCPDF)
+function VAE(zlength::Int, enc::ACD, dec::ACD)
     T = eltype(first(params(enc)))
     μp = NoGradArray(zeros(T, zlength))
     σ2p = NoGradArray(ones(T, zlength))
@@ -61,7 +61,7 @@ Evidence lower boundary of the VAE model. `β` scales the KLD term.
 """
 function elbo(m::AbstractVAE, x::AbstractArray; β=1)
     z = rand(m.encoder, x)
-    llh = mean(loglikelihood(m.decoder, x, z))
+    llh = mean(logpdf(m.decoder, x, z))
     kld = mean(kl_divergence(m.encoder, m.prior, x))
     llh - β*kld
 end

src/utils/flux_ode_decoder.jl

@@ -7,7 +7,7 @@ export FluxODEDecoder
 Uses a Flux `model` as ODE and solves it for the given time span.
 The solver can be conveniently called by (dec::FluxODEDecoder)(z), which
 assumes that all parameters of the neural ODE and its initial conditions are
-passed in as one long vector i.e.: z = vcat(destructure(dec.model), u0).
+passed in as one long vector i.e.: z = vcat(Flux.destructure(dec.model)[1], u0).
 Can use any Flux model as neural ODE. The adjoint is computed via ForwardDiff.
 
 # Arguments
@@ -17,22 +17,26 @@ Can use any Flux model as neural ODE. The adjoint is computed via ForwardDiff.
 * `model`: Flux model
 * `observe`: Observation operator. Function that receives ODESolution and
   outputs the observation. Default: observe(sol) = reshape(hcat(sol.u...),:)
+* `zlength`: length(model) + slength
+* `restructure`: function that maps vector of ODE params to `model`
 """
 mutable struct FluxODEDecoder
     slength::Int
     timesteps::Vector
     model
     observe::Function
-    _zlength::Int
+    zlength::Int
+    restructure::Function
 end
 
 # TODO: FluxODEDecoder{M} fails during training because forward diff wants to
 #       stick dual number in there...
 function FluxODEDecoder(slength::Int, tlength::Int, dt::T,
                         model, observe::Function) where T
     timesteps = range(T(0), step=dt, length=tlength)
-    _zlength = length(destructure(model)) + slength
-    FluxODEDecoder(slength, timesteps, model, observe, _zlength)
+    ps, restructure = Flux.destructure(model)
+    zlength = length(ps) + slength
+    FluxODEDecoder(slength, timesteps, model, observe, zlength, restructure)
 end
 
 function FluxODEDecoder(slength::Int, tlength::Int, dt::Real, model)
@@ -41,11 +45,11 @@ function FluxODEDecoder(slength::Int, tlength::Int, dt::Real, model)
 end
 
 function (dec::FluxODEDecoder)(z::AbstractVector, observe::Function)
-    @assert length(z) == dec._zlength
+    @assert length(z) == dec.zlength
     ps = z[1:end-dec.slength]
     u0 = z[end-dec.slength+1:end]
-    dec.model = restructure(dec.model, ps)
-    z = vcat(destructure(dec.model), u0)
+    dec.model = dec.restructure(ps)
+    z = vcat(Flux.destructure(dec.model)[1], u0)
     tspan = (dec.timesteps[1], dec.timesteps[end])
     dudt_(u::AbstractVector, ps, t) = dec.model(u)
     prob = ODEProblem(dudt_, u0, tspan, ps)

src/utils/utils.jl

@@ -20,7 +20,8 @@ If lastlayer is "linear", then the last layer is forced to be Dense with linear
 
 It is also possible to specify dimensions in a vector.
 """
-layer_builder(k::Vector{Int},l::Vector,f::Vector) = Flux.Chain(map(i -> i[1](i[3],i[4],i[2]),zip(l,f,k[1:end-1],k[2:end]))...)
+layer_builder(k::Vector{Int},l::Vector,f::Vector) =
+    Flux.Chain(map(i -> i[1](i[3],i[4],i[2]),zip(l,f,k[1:end-1],k[2:end]))...)
 
 layer_builder(d::Int,k::Int,o::Int,n::Int, args...) =
     layer_builder(vcat(d,fill(k,n-1)...,o), args...)
@@ -109,37 +110,3 @@ Discriminator loss for true data score `st` and generated data score `sg`.
 """
 discriminator_loss(T,st,sg) = - T(0.5)*(mean(log.(st .+ eps(T))) + mean(log.(1 .- sg) .+ eps(T)))
 discriminator_loss(st,sg) = discriminator_loss(Float32,st,sg)
-
-
-"""
-    destructure(m)
-
-Returns all parameters of a Flux model in one long vector.
-This includes **all** `AbstractArray` fields
-(Adapted from DiffEqFlux.jl)
-"""
-function destructure(m)
-    xs = []
-    Flux.fmap(m) do x
-        x isa AbstractArray && push!(xs, x)
-        return x
-    end
-    return vcat(vec.(xs)...)
-end
-
-"""
-    restructure(m, xs::AbstractVector)
-
-Populate a Flux model with parameters as given in a long vector of xs.
-xs must include **all** `AbstractArray` fields.
-(Adapted from DiffEqFlux.jl)
-"""
-function restructure(m, xs::AbstractVector)
-    i = 0
-    Flux.fmap(m) do x
-        x isa AbstractArray || return x
-        x = reshape(xs[i.+(1:length(x))], size(x))
-        i += length(x)
-        return x
-    end
-end

test/models/gan.jl

@@ -6,15 +6,15 @@
     batchsize = 20
     T = Float32
 
-    test_data = hcat(ones(T,xlen,Int(batchsize/2)), -ones(T,xlen,Int(batchsize/2))) #|> gpu
+    test_data = hcat(ones(T,xlen,Int(batchsize/2)), -ones(T,xlen,Int(batchsize/2))) |> gpu
 
     gen = GenerativeModels.stack_layers([zlen, 10, 10, xlen], relu)
     gen_dist = CMeanGaussian{DiagVar}(gen, NoGradArray(ones(T,xlen)))
 
     disc = GenerativeModels.stack_layers([xlen, 10, 10, 1], relu, Dense; last = Flux.σ)
     disc_dist = CMeanGaussian{DiagVar}(disc, NoGradArray(ones(T,1)))
 
-    model = GAN(zlen, gen_dist, disc_dist) #|> gpu
+    model = GAN(zlen, gen_dist, disc_dist) |> gpu
 
     zs = rand(model.prior, batchsize)
     @test size(zs) == (zlen, batchsize)
@@ -60,7 +60,7 @@
     @test all(param_change(params_disc, model.discriminator))
     @test !any(param_change(params_gen, model.generator))
 
-    msg = @capture_out show(model)
+    msg = sprint(show, model)
     @test occursin("GAN", msg)
 
     Random.seed!()

test/models/rodent.jl

@@ -57,5 +57,7 @@
     # plot!(p, reconstruct(rodent, test_data), ls=:dash)
     # display(p)
 
+    @test occursin("Rodent", sprint(show, rodent))
+
     Random.seed!()  # reset the seed
 end

test/models/vae.jl

@@ -47,15 +47,15 @@
         xlen = 4
         zlen = 2
         batch = 20
-        test_data = hcat(ones(T,xlen,Int(batch/2)), -ones(T,xlen,Int(batch/2))) # |> gpu
+        test_data = hcat(ones(T,xlen,Int(batch/2)), -ones(T,xlen,Int(batch/2))) |> gpu
 
         enc = GenerativeModels.stack_layers([xlen, 10, 10, zlen], relu, Dense)
         enc_dist = CMeanGaussian{DiagVar}(enc, NoGradArray(ones(T,zlen)))
 
         dec = GenerativeModels.stack_layers([zlen, 10, 10, xlen], relu, Dense)
         dec_dist = CMeanGaussian{DiagVar}(dec, NoGradArray(ones(T,xlen)))
 
-        model = VAE(zlen, enc_dist, dec_dist) # |> gpu
+        model = VAE(zlen, enc_dist, dec_dist) |> gpu
 
         # test training
         params_init = get_params(model)
@@ -76,7 +76,7 @@
         @debug maximum(abs.(test_data - xs))
         @test all(abs.(test_data - xs) .< 0.2) 
 
-        msg = @capture_out show(model)
+        msg = sprint(show, model)
         @test occursin("VAE", msg)
         Random.seed!()  # reset the seed
     end

test/runtests.jl

@@ -1,5 +1,4 @@
 using Test
-using Suppressor
 using Logging
 using Parameters
 using Random

test/utils/flux_ode_decoder.jl

@@ -6,7 +6,7 @@
 
     ode = Dense(slength, slength)
     dec = FluxODEDecoder(slength, xlength, dt, ode)
-    ps = GenerativeModels.destructure(ode)
+    ps = Flux.destructure(ode)[1]
     u0 = rand(slength)
     z  = vcat(ps, u0)
 

test/utils/saveload.jl

@@ -37,7 +37,7 @@
 
     opt = ADAM()
     lossf(x) = -elbo(model, x, β=1e-3)
-    data = [(randn(Float32, tlen),)]
+    data = [(randn(Float32, tlen, 1),)]
     Flux.train!(lossf, params(model), data, opt)
     params_trained = get_params(model)
 

test/utils/utils.jl

@@ -32,19 +32,4 @@
     @test m.layers[2] isa Dense
     @test m.layers[2].σ == identity
 
-    @testset "destructure/restructure" begin
-        m = Dense(3,3)
-        σ = NoGradArray(ones(Float32,3))
-        g = CMeanGaussian{DiagVar}(m, σ)
-
-        z = GenerativeModels.destructure(g)
-        @test z isa Vector
-        @test length(z) == 3*3 + 3 + 3
-
-        _z = rand(Float32,3*3 + 3 + 3)
-        _g = GenerativeModels.restructure(g, _z)
-
-        @test _g.mapping.W[1] != g.mapping.W[1]
-        @test _g.mapping.W[1] == _z[1]
-    end
 end