search

TuringLang / TuringCallbacks.jl

http://turinglang.org/TuringCallbacks.jl/
MIT License
13 stars 8 forks source link

TensorBoardCallback doesn't support S3 #40

Closed KDr2 closed 1 year ago

KDr2 commented 1 year ago

Ran the code below:

using Minio, Turing, TuringCallbacks

@model function demo(x)
    s ~ InverseGamma(2, 3)
    m ~ Normal(0, √s)
    for i in eachindex(x)
        x[i] ~ Normal(m, √s)
    end
end

xs = randn(100) .+ 1;
model = demo(xs);

# Number of MCMC samples/steps
num_samples = 10_000
num_adapts = 100

# Sampling algorithm to use
alg = NUTS(num_adapts, 0.65)

# Create the callback

# run a minio server by code
minio_server = Minio.Server(mktempdir(); address="localhost:9000")
run(minio_server, wait=false)
config = MinioConfig("http://localhost:9000")
s3_create_bucket(config, "tensorboard-data")
path = S3Path("s3://tensorboard-data/logdir/"; config=config)

# if your minio server is setup manually, just use the code below
# config = MinioConfig("http://localhost:9000")
# path = S3Path("s3://tensorboard-data/run/"; config=config)
callback = TensorBoardCallback(path)

# Sample
chain = sample(model, alg, num_samples; callback = callback)

I got

┌ Info: Found initial step size                                                                          
└   ϵ = 0.05         
Sampling 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| Time: 0:00:10
ERROR: LoadError: MethodError: no method matching getindex(::Tuple{String}, ::String)                                                                                                                              Closest candidates are:                                                                                                                                                                                            
  getindex(::Tuple, ::Int64) at tuple.jl:29                                                                                                                                                                        
  getindex(::Tuple, ::Integer) at tuple.jl:30                                                            
  getindex(::Tuple, ::Colon) at tuple.jl:33                                                                                                                                                                        
  ...                                                                                                                                                                                                              Stacktrace:                                                                                                                                                                                                        
  [1] getindex(fp::S3Path{MinioConfig}, idx::String)                                                     
    @ FilePathsBase ~/.julia/packages/FilePathsBase/9kSEl/src/path.jl:159
  [2] (::TuringCallbacks.var"#17#18"{Base.Pairs{Symbol, Int64, Tuple{Symbol}, NamedTuple{(:nadapts,), Tuple{Int64}}}, TensorBoardCallback{S3Path{MinioConfig}, Nothing, Nothing, Nothing}, Turing.Inference.HMCTran
sition{NamedTuple{(:s, :m), Tuple{Tuple{Vector{Float64}, Vector{String}}, Tuple{Vector{Float64}, Vector{String}}}}, NamedTuple{(:n_steps, :is_accept, :acceptance_rate, :log_density, :hamiltonian_energy, :hamilto
nian_energy_error, :max_hamiltonian_energy_error, :tree_depth, :numerical_error, :step_size, :nom_step_size), Tuple{Int64, Bool, Float64, Float64, Float64, Float64, Float64, Int64, Bool, Float64, Float64}}, Floa
t64}, Turing.Inference.HMCState{DynamicPPL.TypedVarInfo{NamedTuple{(:s, :m), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:s, Setfield.IdentityLens}, Int64}, Vector{InverseGamma{Float64}}, Vector{AbstractPPL.VarName{:s, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:m, Setfield.IdentityLens}, Int64}, Vector{Normal{Float64}}, Vector{AbstractPPL.VarName{:m, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, AdvancedHMC.HMCKernel{AdvancedHMC.FullMomentumRefreshment, AdvancedHMC.Trajectory{AdvancedHMC.Multinomia
lTS, AdvancedHMC.Leapfrog{Float64}, AdvancedHMC.GeneralisedNoUTurn{Float64}}}, AdvancedHMC.Hamiltonian{AdvancedHMC.DiagEuclideanMetric{Float64, Vector{Float64}}, AdvancedHMC.GaussianKinetic, Base.Fix1{typeof(Log
DensityProblems.logdensity), LogDensityProblemsAD.LogDensityProblemsADForwardDiffExt.ForwardDiffLogDensity{LogDensityFunction{DynamicPPL.TypedVarInfo{NamedTuple{(:s, :m), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:s, Setfield.IdentityLens}, Int64}, Vector{InverseGamma{Float64}}, Vector{AbstractPPL.VarName{:s, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:m, Setfield.IdentityLens}, Int64}, Vector{Normal{Float64}}, Vector{AbstractPPL.VarName{:m, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, DynamicPPL
.Model{typeof(demo), (:x,), (), (), Tuple{Vector{Float64}}, Tuple{}, DynamicPPL.DefaultContext}, DynamicPPL.SamplingContext{DynamicPPL.Sampler{NUTS{Turing.Essential.ForwardDiffAD{0}, (), AdvancedHMC.DiagEuclidea
nMetric}}, DynamicPPL.DefaultContext, Random._GLOBAL_RNG}}, ForwardDiff.Chunk{2}, ForwardDiff.Tag{Turing.TuringTag, Float64}, ForwardDiff.GradientConfig{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 2, Ve
ctor{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 2}}}}}, Turing.Inference.var"#∂logπ∂θ#44"{LogDensityProblemsAD.LogDensityProblemsADForwardDiffExt.ForwardDiffLogDensity{LogDensityFuncti
on{DynamicPPL.TypedVarInfo{NamedTuple{(:s, :m), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:s, Setfield.IdentityLens}, Int64}, Vector{InverseGamma{Float64}}, Vector{AbstractPPL.VarName{:s, Setfield.Ident
ityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:m, Setfield.IdentityLens}, Int64}, Vector{Normal{Float64}}, Vector{AbstractPPL.VarName{:m, Setfield.Id
entityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, DynamicPPL.Model{typeof(demo), (:x,), (), (), Tuple{Vector{Float64}}, Tuple{}, DynamicPPL.DefaultContext}, DynamicPPL.SamplingContext
{DynamicPPL.Sampler{NUTS{Turing.Essential.ForwardDiffAD{0}, (), AdvancedHMC.DiagEuclideanMetric}}, DynamicPPL.DefaultContext, Random._GLOBAL_RNG}}, ForwardDiff.Chunk{2}, ForwardDiff.Tag{Turing.TuringTag, Float64
}, ForwardDiff.GradientConfig{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 2, Vector{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 2}}}}}}, AdvancedHMC.PhasePoint{Vector{Float64},
 AdvancedHMC.DualValue{Float64, Vector{Float64}}}, AdvancedHMC.Adaptation.StanHMCAdaptor{AdvancedHMC.Adaptation.WelfordVar{Float64, Vector{Float64}}, AdvancedHMC.Adaptation.NesterovDualAveraging{Float64}}}, Base
.Fix1{typeof(TuringCallbacks.filter_param_and_value), TensorBoardCallback{S3Path{MinioConfig}, Nothing, Nothing, Nothing}}, TensorBoardLogger.TBLogger{String, IOStream}, S3Path{MinioConfig}})()
    @ TuringCallbacks ~/.julia/packages/TuringCallbacks/5xDTB/src/callbacks/tensorboard.jl:173
  [3] with_logstate(f::Function, logstate::Any)
    @ Base.CoreLogging ./logging.jl:511
  [4] with_logger
    @ ./logging.jl:623 [inlined]
  [5] #_#16
    @ ~/.julia/packages/TuringCallbacks/5xDTB/src/callbacks/tensorboard.jl:171 [inlined]
  [6] macro expansion
    @ ~/.julia/packages/AbstractMCMC/F9Hbk/src/sample.jl:140 [inlined]
  [7] macro expansion
    @ ~/.julia/packages/ProgressLogging/6KXlp/src/ProgressLogging.jl:328 [inlined]
  [8] (::AbstractMCMC.var"#20#21"{Bool, String, TensorBoardCallback{S3Path{MinioConfig}, Nothing, Nothing, Nothing}, Int64, Int64, Base.Pairs{Symbol, Int64, Tuple{Symbol}, NamedTuple{(:nadapts,), Tuple{Int64}}},
 Random.TaskLocalRNG, DynamicPPL.Model{typeof(demo), (:x,), (), (), Tuple{Vector{Float64}}, Tuple{}, DynamicPPL.DefaultContext}, DynamicPPL.Sampler{NUTS{Turing.Essential.ForwardDiffAD{0}, (), AdvancedHMC.DiagEuc
lideanMetric}}, Int64, Int64})()
    @ AbstractMCMC ~/.julia/packages/AbstractMCMC/F9Hbk/src/logging.jl:12
  [9] with_logstate(f::Function, logstate::Any)
    @ Base.CoreLogging ./logging.jl:511
 [10] with_logger(f::Function, logger::LoggingExtras.TeeLogger{Tuple{LoggingExtras.EarlyFilteredLogger{TerminalLoggers.TerminalLogger, AbstractMCMC.var"#1#3"{Module}}, LoggingExtras.EarlyFilteredLogger{Logging.C
onsoleLogger, AbstractMCMC.var"#2#4"{Module}}}})
    @ Base.CoreLogging ./logging.jl:623
 [11] with_progresslogger(f::Function, _module::Module, logger::Logging.ConsoleLogger)
    @ AbstractMCMC ~/.julia/packages/AbstractMCMC/F9Hbk/src/logging.jl:36
 [12] macro expansion
    @ ~/.julia/packages/AbstractMCMC/F9Hbk/src/logging.jl:11 [inlined]
 [13] mcmcsample(rng::Random.TaskLocalRNG, model::DynamicPPL.Model{typeof(demo), (:x,), (), (), Tuple{Vector{Float64}}, Tuple{}, DynamicPPL.DefaultContext}, sampler::DynamicPPL.Sampler{NUTS{Turing.Essential.Forw
ardDiffAD{0}, (), AdvancedHMC.DiagEuclideanMetric}}, N::Int64; progress::Bool, progressname::String, callback::TensorBoardCallback{S3Path{MinioConfig}, Nothing, Nothing, Nothing}, discard_initial::Int64, thinnin
g::Int64, chain_type::Type, kwargs::Base.Pairs{Symbol, Int64, Tuple{Symbol}, NamedTuple{(:nadapts,), Tuple{Int64}}})
    @ AbstractMCMC ~/.julia/packages/AbstractMCMC/F9Hbk/src/sample.jl:116
 [14] sample(rng::Random.TaskLocalRNG, model::DynamicPPL.Model{typeof(demo), (:x,), (), (), Tuple{Vector{Float64}}, Tuple{}, DynamicPPL.DefaultContext}, sampler::DynamicPPL.Sampler{NUTS{Turing.Essential.ForwardD
iffAD{0}, (), AdvancedHMC.DiagEuclideanMetric}}, N::Int64; chain_type::Type, resume_from::Nothing, progress::Bool, nadapts::Int64, discard_adapt::Bool, discard_initial::Int64, kwargs::Base.Pairs{Symbol, TensorBo
ardCallback{S3Path{MinioConfig}, Nothing, Nothing, Nothing}, Tuple{Symbol}, NamedTuple{(:callback,), Tuple{TensorBoardCallback{S3Path{MinioConfig}, Nothing, Nothing, Nothing}}}})
    @ Turing.Inference ~/.julia/packages/Turing/Suzsv/src/inference/hmc.jl:133
 [15] #sample#2
    @ ~/.julia/packages/Turing/Suzsv/src/inference/Inference.jl:146 [inlined]
 [16] #sample#1
    @ ~/.julia/packages/Turing/Suzsv/src/inference/Inference.jl:136 [inlined]
 [17] top-level scope
    @ cb.jl:27

I noticed that TensorBoardLogger.jl support S3 and Minio, and copied the code from its test cases (https://github.com/JuliaLogging/TensorBoardLogger.jl/blob/875fdb47be6dd02a613d0d160b579faa62788715/test/runtests.jl#L17).

Is it possible to make TensorBoardCallback support S3?

Thanks.

torfjelde commented 1 year ago

Have you tried just constructing the TBLogger directly and passing that?

KDr2 commented 1 year ago

constructing the TBLogger directly and passing that

Thanks, this works, but it's extremely slow, it took hours to finish the above case... :(

torfjelde commented 1 year ago

Okay, that's very strange. But it doesn't take hours to finish with out S3, right?

But if it takes hours to finish once you use S3-integration, that sounds more like an issue with either TensorBoardLogger.jl or Minio.jl :confused: If you give us the TBLogger, then all we do is just call logging statements. Maybe try benchmarking a function with logging to S3 and one logging locally?

KDr2 commented 1 year ago

that sounds more like an issue with either TensorBoardLogger.jl or Minio.jl

I think you are right.

When num_sample=100:

When num_sample=1000:

torfjelde commented 1 year ago

Furthermore, given that there's no specialization for Minio.jl in TensorBoardLogger.jl, I'm guessing it's a Minio.jl issue.