Merge pull request #961 from JuliaRobotics/feat/4Q20/fluxnn

add FluxModelsDistribution

Author: dehann

Project.toml

@@ -6,6 +6,7 @@ version = "0.16.0"
 
 [deps]
 ApproxManifoldProducts = "9bbbb610-88a1-53cd-9763-118ce10c1f89"
+Base64 = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
@@ -39,7 +40,7 @@ UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 [compat]
 ApproxManifoldProducts = "0.1, 0.2"
 Combinatorics = "1.0"
-DistributedFactorGraphs = "0.10.5"
+DistributedFactorGraphs = "0.10.6"
 Distributions = "0.20, 0.21, 0.22, 0.23, 0.24"
 DocStringExtensions = "0.8, 0.9, 0.10, 1"
 FileIO = "1.0.2, 1.1, 1.2"
@@ -60,8 +61,10 @@ TimeZones = "1.3.1"
 julia = "1.4"
 
 [extras]
+BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0"
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test"]
+test = ["BSON", "Flux", "Test"]

src/BeliefTypes.jl

@@ -25,7 +25,7 @@ struct NonparametricMessage <: MessageType end
 struct ParametricMessage <: MessageType end
 
 
-const SamplableBelief = Union{Distributions.Distribution, KernelDensityEstimate.BallTreeDensity, AliasingScalarSampler}
+const SamplableBelief = Union{Distributions.Distribution, KernelDensityEstimate.BallTreeDensity, AliasingScalarSampler, FluxModelsDistribution}
 
 abstract type PackedSamplableBelief end
 

src/Deprecated.jl

@@ -858,6 +858,8 @@ function initVariable!(fgl::AbstractDFG,
 end
 
 
+export uppA
+
 #global pidx
 global pidx = 1
 global pidl = 1

src/FactorGraph.jl

@@ -558,6 +558,15 @@ end
 
 # import IncrementalInference: prepgenericconvolution, convert
 
+"""
+    $SIGNATURES
+
+Function to calculate measurement dimension from factor sampling.
+
+Notes
+- Will not work in all situations, but good enough so far.
+  - # TODO standardize
+"""
 function calcZDim(usrfnc::T, Xi::Vector{<:DFGVariable})::Int where {T <: FunctorInferenceType}
   # zdim = T != GenericMarginal ? size(getSample(usrfnc, 2)[1],1) : 0
   zdim = if T != GenericMarginal

src/Factors/DefaultPrior.jl

@@ -52,10 +52,10 @@ mutable struct PackedPrior <: PackedInferenceType
   PackedPrior(z::AS) where {AS <: AbstractString} = new(z)
 end
 function convert(::Type{PackedPrior}, d::Prior)
-  PackedPrior(string(d.Z))
+  PackedPrior(convert(PackedSamplableBelief, d.Z))
 end
 function convert(::Type{Prior}, d::PackedPrior)
-  Prior(extractdistribution(d.Z))
+  Prior(convert(SamplableBelief, d.Z))
 end
 
 

src/Factors/Mixture.jl

@@ -53,10 +53,10 @@ end
 
 # should not be called in case of Prior
 (s::Mixture)( res::AbstractArray{<:Real},
-              userdata::FactorMetadata,
+              fmd::FactorMetadata,
               idx::Int,
               meas::Tuple,
-              X... ) = s.mechanics(res, userdata, idx, meas, X...)
+              X... ) = s.mechanics(res, fmd, idx, meas, X...)
 #
 
 

src/Flux/FluxModelsDistribution.jl

@@ -0,0 +1,61 @@
+
+
+@info "IncrementalInference is adding Flux related functionality."
+
+
+
+# the factor definitions
+export FluxModelsDistribution
+# some utilities
+
+
+# Required packages
+using .Flux
+
+using Random, Statistics
+# using DistributedFactorGraphs
+
+# import Base: convert
+import Random: rand
+
+
+
+function rand(nfb::FluxModelsDistribution,
+              N::Int=1 )
+  #
+
+  # number of predictors to choose from, and choose random subset
+  numModels = length(nfb.models)
+  allPreds = 1:numModels |> collect # 1:Npreds |> collect
+  # TODO -- compensate when there arent enough prediction models
+  if numModels < N
+    reps = (N ÷ numModels) + 1
+    allPreds = repeat(allPreds, reps  )
+    resize!(allPreds, N)
+  end
+  # samples for the order in which to use models, dont shuffle if N models
+  # can suppress shuffle for NN training purposes
+  1 < numModels && nfb.shuffle[] ? shuffle!(allPreds) : nothing
+
+  # generate the measurements
+  meas = zeros(nfb.outputDim..., N)
+  for i in 1:N
+    meas[:,i] = (nfb.models[allPreds[i]])(nfb.data)
+  end
+
+  return meas
+end
+
+
+FluxModelsDistribution( inDim::NTuple{ID,Int}, 
+                        outDim::NTuple{OD,Int}, 
+                        models::Vector{P}, 
+                        data::D,
+                        shuffle::Bool=true,
+                        serializeHollow::Bool=false ) where {ID,OD,P,D<:AbstractArray} = FluxModelsDistribution{ID,OD,P,D}(inDim, outDim, models, data, Ref(shuffle), Ref(serializeHollow) )
+#
+
+
+
+
+#

src/Flux/FluxModelsSerialization.jl

@@ -0,0 +1,134 @@
+# Serialization functions for Flux models that depend on BSON
+
+@info "IncrementalInference is adding Flux/BSON serialization functionality."
+
+export PackedFluxModelsDistribution
+
+using .BSON
+using Base64
+
+import Base: convert
+
+
+mutable struct PackedFluxModelsDistribution <: IIF.PackedSamplableBelief
+  # shape of the input data
+  inputDim::Vector{Int}
+  # shape of the output data
+  outputDim::Vector{Int}
+  # actual Flux models (Base64 encoded binary)
+  mimeTypeModel::String
+  models::Vector{String}
+  # the data used for prediction, must be <: AbstractArray
+  mimeTypeData::String
+  data::String
+  # shuffle model predictions relative to particle index at each sampling
+  shuffle::Bool
+  # false for default serialization with model info, set true for separate storage of models 
+  serializeHollow::Bool
+    # TODO remove requirement and standardize sampler API
+    # specialSampler::Symbol
+  # field name usage to direct the IIF serialization towards JSON method
+  PackedSamplableTypeJSON::String
+end
+
+function _serializeFluxModelBase64(model::Flux.Chain)
+  io = IOBuffer()
+  iob64 = Base64EncodePipe(io)
+  BSON.@save iob64 model
+  close(iob64)
+  return String(take!(io))
+end
+
+function _deserializeFluxModelBase64(smodel::AbstractString)
+  iob64 = PipeBuffer(base64decode(smodel))
+  BSON.@load iob64 model
+  close(iob64)
+  return model
+end
+
+function _serializeFluxDataBase64(data::AbstractArray)
+  io = IOBuffer()
+  iob64 = Base64EncodePipe(io)
+  BSON.@save iob64 data
+  close(iob64)
+  return String(take!(io))
+end
+
+function _deserializeFluxDataBase64(sdata::AbstractString)
+  iob64 = PipeBuffer(base64decode(sdata))
+  BSON.@load iob64 data
+  close(iob64)
+  return data
+end
+
+
+function convert( ::Union{Type{<:PackedSamplableBelief},Type{<:PackedFluxModelsDistribution}}, 
+                  obj::FluxModelsDistribution)
+  #
+
+    # and the specialSampler function -- likely to be deprecated
+    # specialSampler = Symbol(obj.specialSampler)
+  # fields to persist
+  inputDim = collect(obj.inputDim)
+  outputDim = collect(obj.outputDim)
+  models = Vector{String}()
+  # store all models as Base64 Strings (using BSON)
+  if !obj.serializeHollow[]
+    resize!(models, length(obj.models))
+    # serialize the Vector of Flux models (each one individually)
+    models .= _serializeFluxModelBase64.(obj.models)
+    # also store data as Base64 String, using BSON
+    sdata = _serializeFluxDataBase64(obj.data)
+    mimeTypeData = "application/bson/octet-stream/base64"
+  else
+    # store one just model to preserve the type (allows resizing on immutable Ref after deserialize)
+    push!(models,_serializeFluxModelBase64(obj.models[1]))
+    # at least capture the type of how the data looks for future deserialization
+    sdata = string(typeof(obj.data))
+    mimeTypeData = "application/text"
+  end
+  mimeTypeModel = "application/bson/octet-stream/base64"
+  
+  # and build the JSON-able object
+  packed = PackedFluxModelsDistribution(inputDim, 
+                                        outputDim, 
+                                        mimeTypeModel,
+                                        models, 
+                                        mimeTypeData,
+                                        sdata, 
+                                        obj.shuffle[], 
+                                        obj.serializeHollow[], 
+                                        # specialSampler,
+                                        "IncrementalInference.PackedFluxModelsDistribution" )
+  #
+  return JSON2.write(packed)
+end
+
+
+
+function convert( ::Union{Type{<:SamplableBelief},Type{FluxModelsDistribution}}, 
+                  obj::PackedFluxModelsDistribution)
+  #
+
+  obj.serializeHollow && @warn("Deserialization of FluxModelsDistribution.serializationHollow=true is not yet well developed, please open issues at IncrementalInference.jl accordingly.")
+
+  # specialSampler likely to be deprecated
+  # specialSampler = getfield(Main, obj.specialSampler)
+  
+  # deserialize
+  # @assert obj.mimeTypeModel == "application/bson/octet-stream/base64"
+  models = _deserializeFluxModelBase64.(obj.models)
+  
+  # @assert obj.mimeTypeData == "application/bson/octet-stream/base64"
+  data = !obj.serializeHollow ? _deserializeFluxDataBase64.(obj.data) : zeros(0)
+
+  return FluxModelsDistribution((obj.inputDim...,), 
+                                (obj.outputDim...,), 
+                                models, 
+                                data, 
+                                obj.shuffle, 
+                                obj.serializeHollow  )
+end
+
+
+#

src/Flux/entities.jl

@@ -0,0 +1,19 @@
+# entities immediately available as private members in IIF, but require Flux for actual use
+
+
+struct FluxModelsDistribution{ID,OD,P,D<:AbstractArray}
+  # shape of the input data
+  inputDim::NTuple{ID,Int}
+  # shape of the output data
+  outputDim::NTuple{OD,Int}
+  # actual Flux models
+  models::Vector{P}
+  # the data used for prediction, must be <: AbstractArray
+  data::D
+  # shuffle model predictions relative to particle index at each sampling
+  shuffle::Base.RefValue{Bool}
+  # false for default serialization with model info, set true for separate storage of models 
+  serializeHollow::Base.RefValue{Bool}
+    # # TODO remove requirement and standardize sampler API
+    # specialSampler::Function
+end

src/IncrementalInference.jl

@@ -84,6 +84,7 @@ export AbstractDFG,
   listDataEntries,
   FolderStore,
   addBlobStore!,
+  addData!,
   getData,
   DFGVariable,
   DFGVariableSummary, 
@@ -96,6 +97,9 @@ export AbstractDFG,
 export FunctorInferenceType, PackedInferenceType
 export AbstractPrior, AbstractRelativeFactor, AbstractRelativeFactorMinimize
 
+# not sure if this is necessary
+export convert
+
 export *,
   notifyCSMCondition,
   CSMHistory,
@@ -367,10 +371,6 @@ export *,
   GenericMarginal,
   PackedGenericMarginal,
 
-  uppA,
-  convert,
-  extractdistribution,
-
   # factor graph operating system utils (fgos)
   saveTree,
   loadTree,
@@ -422,6 +422,7 @@ getFactorOperationalMemoryType(dfg::SolverParams) = CommonConvWrapper
 
 include("AliasScalarSampling.jl")
 include("CliqueTypes.jl")
+include("Flux/entities.jl")
 include("BeliefTypes.jl")
 include("JunctionTreeTypes.jl")
 include("FactorGraph.jl")
@@ -492,9 +493,15 @@ include("Deprecated.jl")
 
 exportimg(pl) = error("Please do `using Gadfly` before IncrementalInference is used to allow image export.")
 function __init__()
-    @require InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" include("RequireInteractiveUtils.jl")
+  @require InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" include("RequireInteractiveUtils.jl")
+
+  @require Gadfly="c91e804a-d5a3-530f-b6f0-dfbca275c004" include("EmbeddedPlottingUtils.jl")
 
-    @require Gadfly="c91e804a-d5a3-530f-b6f0-dfbca275c004" include("EmbeddedPlottingUtils.jl")
+  # combining neural networks natively into the non-Gaussian  factor graph object
+  @require Flux="587475ba-b771-5e3f-ad9e-33799f191a9c" begin
+    include("Flux/FluxModelsDistribution.jl")
+    @require BSON="fbb218c0-5317-5bc6-957e-2ee96dd4b1f0" include("Flux/FluxModelsSerialization.jl")
+  end
 end
 
 # Old code that might be used again