weak_learner_augmentation.jl

# include(Pkg.dir("ExperimentalAnalysis","src","ExperimentalAnalysis.jl"))
using TensorFlow, DeterministicPolicyGradient, ValueHistories, JLD, DataFrames, BasisFunctionExpansions
using Plots
default(size=(1600,900), show=false, legend=false, lab="")
const input_size  = 1
const output_size = 1
const neurons     = [2000,2000]
const dir         = "weak_learner_experiments"
const N_burnin    = 200
const FloatType   = Float64

function fanin(T,dims...)
    s = √(6/(sum(dims)))
    2s*rand(T,dims...)-s
end

function fanin2(T,dims...)
    s = √(2/(dims[1]))
    s*randn(T,dims...)
end

function runstuff(stepsize=2e-4,weight_cost=0.000001,plotflag=true,verbose=false)

    session      = Session(Graph())

    # TODO: check if these must be moved outside function because they cause allocations
    x_ = placeholder(FloatType, shape =[-1,input_size])
    y_ = placeholder(FloatType, shape =[-1,output_size])
    keep_prob = placeholder(FloatType)
    W1 = Variable(fanin2(FloatType, input_size, neurons[1]), name  ="weights1")
    W2 = Variable(fanin2(FloatType, neurons[1], neurons[2]), name  ="weights2")
    W3 = Variable(fanin2(FloatType, neurons[2], output_size), name ="weights3")
    B1 = Variable(0.002*rand(FloatType,neurons[1])-0.001, name             ="bias1")
    B2 = Variable(0.002*rand(FloatType,neurons[2])-0.001, name             ="bias2")
    B3 = Variable(0*ones(FloatType,output_size), name                      ="bias3")

    dropout(x) = nn.dropout(x,keep_prob)
    l1   = x_*W1 + B1 |> nn.relu |> dropout
    l2   = l1*W2 + B2 |> nn.relu |> dropout
    q    = l2*W3 + B3


    loss         = reduce_mean((y_ - q).^2)
    weight_decay = weight_cost*(reduce_sum(W1.^2) + reduce_sum(W2.^2) + reduce_sum(W3.^2))

    cost         = loss + weight_decay
    train_step   = train.minimize(train.AdamOptimizer(stepsize), cost)

    fun(x)       = (x + x^2 - 5)*(x < 3) - (x )*(x >= 3)

    N_initial = 50
    initial_x = 6rand(N_initial)-3 |> sort
    initial_x = initial_x''
    initial_y = fun.(initial_x)

    N_second  = 100
    second_x  = 9rand(N_second)-3 |> sort
    second_x  = second_x''
    second_y  = fun.(second_x)

    if plotflag
        gr()
        fig = scatter(second_x, second_y, layout=(1,4))
    end
    run(session, initialize_all_variables())
    vh = History(FloatType)
    vh2 = History(FloatType)

    # Train using initial data distribution
    i = 0
    _,_,Θ = fit_weak(initial_x,initial_y,0)
    for i = 1:N_burnin
        run(session, [train_step], Dict(keep_prob=>0.9, y_ => (initial_y-eval_weak(initial_x,Θ)), x_ => initial_x))
        yh,l,L1,L2 = run(session, [q,loss,l1,l2], Dict(keep_prob=>1., y_ => (initial_y-eval_weak(initial_x,Θ)), x_ => initial_x))
        yh2,l,Θ = fit_weak(initial_x,initial_y,yh)
        push!(vh,i,l)
        verbose && println("Epoch: ", i, " loss: ", l)
        plotflag || continue
        scatter!(second_x, second_y,c=:blue, subplot=1)
        scatter!(initial_x,[yh+yh2 yh yh2],c=[:red :green :magenta],subplot=1)
        scatter!(vh,yscale=:log10,subplot=2)
        update_plot!(fig[1], max_history=4)
        update_plot!(fig[2], max_history=1)
        if i % 5 == 0
            heatmap!(L1, subplot=3)
            heatmap!(L2, subplot=4)
            update_plot!(fig[3], max_history=1)
            update_plot!(fig[4], max_history=1)
        end
        # sleep(0.01)
        gui(fig)
    end

    # Train using modified data distribution
    for i in i+1:i+1000
        run(session, [train_step], Dict(keep_prob=>0.9, y_ => (second_y-eval_weak(second_x,Θ)), x_ => second_x))
        yh,second_l ,L1,L2= run(session, [q,loss,l1,l2], Dict(keep_prob=>1., y_ => (second_y-eval_weak(second_x,Θ)), x_ => second_x))
        yh2,second_l,Θ = fit_weak(second_x,second_y,yh)

        initial_l = run(session, loss, Dict(keep_prob=>1., y_ => (initial_y-eval_weak(initial_x,Θ)), x_ => initial_x))
        verbose && println("Epoch: ", i, " loss: ", initial_l)
        push!(vh,i,initial_l)
        push!(vh2,i,second_l)
        plotflag || continue
        scatter!(second_x, second_y, c=:blue, subplot=1)
        scatter!(second_x,[yh+yh2 yh yh2],c=[:red :green :magenta],subplot=1)
        scatter!(vh,yscale=:log10,subplot=2)
        scatter!(vh2,yscale=:log10,subplot=2)
        update_plot!(fig[1], max_history=4)
        update_plot!(fig[2], max_history=2)
        if i % 5 == 0
            heatmap!(L1, subplot=3)
            heatmap!(L2, subplot=4)
            update_plot!(fig[3], max_history=1)
            update_plot!(fig[4], max_history=1)
        end
        gui()
    end
    @save(joinpath(dir,dir*string(length(readdir(dir)))*".jld"),stepsize,vh.values,vh2.values,weight_cost)
    vh,vh2
end




function runmany(howmany)

    stepsizes    = logspace(-2.5,-1,20)
    weight_costs = logspace(-7,1,20)
    for i = 1:howmany
        stepsize    = stepsizes[rand(1:length(stepsizes))]
        weight_cost = weight_costs[rand(1:length(weight_costs))]
        runstuff(stepsize,weight_cost,false)
    end
end



# # # Polynomial regression
# reg(x) = [ones(x) x x.^2]
# function fit_weak(x,y,yh)
#     e        = y-yh
#     A        = reg(x)
#     n_params = size(A,2)
#     Θ        = [A; 10*eye(n_params)]\[e;zeros(n_params)]
#     yh2      = A*Θ
#     loss     = mean((e-yh2).^2)
#     yh2,loss,Θ
# end
# eval_weak(x,Θ) = reg(x)*Θ





const bfe = UniformRBFE(-3:6, 10)
eval_weak(x,bfa) = bfa(vec(x))
function fit_weak(x,y,yh)
    e    = y-yh |>vec
    vx = vec(x)
    bfa  = BasisFunctionApproximation(e,vx,bfe,100)
    yh2  = eval_weak(vx,bfa)
    loss = mean((e.-yh2).^2)
    yh2,loss,bfa
end