Open deoxyribose opened 3 years ago
I think the issue might be that pick_random_leaf
is not sampling a single random choice, but is itself a generative function that samples multiple random choices. If that's the case, you can either make pick_random_leaf
into a Gen.Distribution
(https://www.gen.dev/dev/ref/extending/#custom_distributions-1) or you can extend the transform with a call to a recursive transform (that you call with @tcall
, see https://www.gen.dev/dev/ref/trace_translators/#Trace-Transform-DSL-1) that walks the tree and populates values of all the choices made by pick_random_leaf
(see https://github.com/probcomp/GenExamples.jl/blob/main/gp_structure/involution_mh.jl#L91-L155 for an example of that approach).
Note that while you can @copy
entire choice maps from one trace to another, you can only @read
and @write
individual random choices.
Thank you for the answer! I'm not sure if it's exactly what you meant, but I tried the recursive transform approach like this:
@gen function pick_random_leaf(node::Node, cur::Int, depth::Int)
if isa(node, LeafNode)
(cur, depth)
elseif @trace(bernoulli(0.5), :recurse_a => cur)
@trace(pick_random_leaf(node.a, get_child(cur, 1, 2), depth+1))
else
@trace(pick_random_leaf(node.b, get_child(cur, 2, 2), depth+1))
end
end
@gen function pick_random_leaf_parent(node::Node, cur::Int, depth::Int)
if isa(node.a, LeafNode) && isa(node.b, LeafNode)
return (cur, depth)
else
if isa(node.a, BranchNode) && isa(node.b, BranchNode)
recurse_a_prob = 0.5
elseif isa(node.a, BranchNode) && !isa(node.b, BranchNode)
recurse_a_prob = 1.
elseif isa(node.b, BranchNode) && !isa(node.a, BranchNode)
recurse_a_prob = 0.
else
error(node)
end
end
if @trace(bernoulli(recurse_a_prob), :recurse_a => cur)
@trace(pick_random_leaf_parent(node.a, get_child(cur, 1, 2), depth+1))
else
@trace(pick_random_leaf_parent(node.b, get_child(cur, 2, 2), depth+1))
end
end
@gen function split_merge_proposal(prev_trace)
tree = prev_trace[:tree]
n_leaf_nodes = count_leaves(tree)
random_split = @trace(bernoulli(0.5), :split)
split = (n_leaf_nodes == 1) ? true : random_split
if split
# select random leaf for splitting
leaf_path = @trace(pick_random_leaf(tree, 1, 0), :leaf_path)
new_node = @trace(make_branch(), :new_node)
else
# select random leaf that will merge with sibling
# sibling needs to be a leaf as well
leaf_path = @trace(pick_random_leaf_parent(tree, 1, 0), :leaf_path)
new_node = @trace(make_leaf(), :new_node)
end
(leaf_path, new_node)
end
@transform walk_tree(cur::Int, leaf_path::Array{Symbol,1}) (model_in, aux_in) to (model_out, aux_out) begin
(leaf_number, leaf_depth) = @read(aux_in[:leaf_path], :discrete)
if leaf_number == cur
new_node_addr = isempty(leaf_path) ? :tree : (:tree => foldr(=>, leaf_path))
@copy(aux_in[:new_node], model_out[new_node_addr])
@copy(model_in[new_node_addr], aux_out[:new_node])
else
recurse_a = @read(aux_in[:leaf_path => :recurse_a => cur], :discrete)
if recurse_a
push!(leaf_path, :a)
@tcall(walk_tree(get_child(cur, 1, 2), leaf_path))
else
push!(leaf_path, :b)
@tcall(walk_tree(get_child(cur, 2, 2), leaf_path))
end
end
end
@transform split_merge_transform (model_in, aux_in) to (model_out, aux_out) begin
(leaf_number, leaf_depth) = @read(aux_in[:leaf_path], :discrete)
new_node = @read(aux_in[:new_node], :discrete)
tree = @read(model_in[:tree], :discrete)
n_leaf_nodes = count_leaves(tree)
random_split = @read(aux_in[:split], :discrete)
split = (n_leaf_nodes == 1) ? true : random_split
@copy(aux_in[:leaf_path], aux_out[:leaf_path])
@write(aux_out[:split], !random_split, :discrete)
@tcall(walk_tree(1,Symbol[]))
end
It does pass the round trip checks, and usually gets many iterations in, but eventually crashes on
ERROR: Did not visit all constraints
Any idea what might be going wrong, or how I can inspect what constraints aren't visited and why?
I've also been trying to do map_optimize on a selection of the continuous choices in the trace, like this:
function select_continuous(node)
selection = DynamicSelection()
leaves = get_path_to_leaf(node)
branches = []
for leaf in leaves
leaf_address = isempty(leaf) ? :tree : (:tree => foldr(=>, leaf))
push!(selection, leaf_address => :mean)
push!(selection, leaf_address => :variance)
for (i,choice) in enumerate(leaf)
if leaf[1:i] ∉ branches
push!(branches, leaf[1:i])
branch_address = isempty(leaf[1:i]) ? :tree : (:tree => foldr(=>, leaf[1:i]))
push!(selection, branch_address => :frac)
end
end
end
return selection
end
function do_inference(model, img, n_iter)
# condition on image
observations = choicemap()
observations[:img] = [img...]
# generate initial trace consistent with observed data
(trace, _) = generate(model, (size(img),), observations)
continous_variables = select_continuous(trace[:tree])
# do MCMC
for iter=1:n_iter
# do MH move on the subtree
trace = replace_split_merge_move(trace)
# optimize continuous variables
trace = map_optimize(trace, continuous_variables)
end
return trace
end;
but this runs into lots of different bugs like
ERROR: ArgumentError: Converting an instance of ReverseDiff.TrackedReal{Float64,Float64,Nothing} to Float64 is not defined. Please use
ReverseDiff.valueinstead.
and
ERROR: AssertionError: length(arr) >= start_idx
which makes me think there's something fundamentally wrong with how I'm setting this up.
In any case, thank you for your time!
Hi,
I'm trying to do involutive MCMC for a 2d non-parametric change-point model (the same one as in #388). The model samples a tree structure that recursively divides an image into rectangles, and then samples pixel values for the whole image given the structure. Model code here, the relevant bits are
I wrote an involution inspired by the GP structure search in GenExamples.jl, where the proposal picks a random node in the tree, and samples a subtree that replaces whatever was at that node. Full code here.
This runs without errors, but few traces are accepted, probably because randomly proposed trees aren't a good fit, and even a well-fitting tree usually has bad values (:mean and :variance) at the leaf nodes.
So I want to write a split-merge involution which proposes to split a leaf node, or merge sibling leaf nodes, along with :mean and :variance parameters derived from the current trace, similar to the mixture example here. To begin with, I skip the parameter proposals and just propose a random split/merge, like this:
Full code here
Just like :choose_subtree_root in the subtree involution, to perform the correct reversal, aux_out needs the same value at :leaf_path that aux_in had, so I did exactly the same
@copy(aux_in[:leaf_path], aux_out[:leaf_path])
But here, it runs for a few iterations and then crashes onERROR: transform round trip check failed
because the copied :leaf_path would, for some reason, differ from the original one, and so the new node that was proposed was spliced back at the wrong address, resulting in different model traces before and after the round trip. I tried with@write(aux_out[:leaf_path], leaf_path, :discrete)
instead, and now I getERROR: KeyError: key :leaf_path not found
for eitherleaf_path = @trace(pick_random_leaf(tree), :leaf_path)
orleaf_path = @trace(pick_random_leaf_parent(tree), :leaf_path)
in the proposal.I hope this makes sense, let me know if I should clarify anything. Ultimately I want to do split-merge jumps that derive the leaf node values from the previous state (e.g. take the means of the mean and variance parameters of two merging leaf nodes), and then do MAP optimization on all the continuous parameters. But maybe there's a better way to do inference?
Thanks in advance for any help.