Learning of Cost Weights

written by Benjamin Gallusser, edited by Jan Funke

Download this page as a Jupyter notebook

motile supports learning of cost weights, given (sparse) ground-truth annotations on the track graph. This tutorial will show you:

What those weights are and how to modify them by hand.
How to annotate nodes and edges as being part of the ground-truth.
How to learn optimal weights for the costs for a given annotated track graph.

Cost Weights

Consider the following example track graph, where each node is a potential object and each edge a potential link of objects between frames:

(click here to see how to create the example track graph)

import motile
import networkx as nx

cells = [
        {"id": 0, "t": 0, "x": 1, "score": 0.8},
        {"id": 1, "t": 0, "x": 25, "score": 0.1},
        {"id": 2, "t": 1, "x": 0, "score": 0.3},
        {"id": 3, "t": 1, "x": 26, "score": 0.4},
        {"id": 4, "t": 2, "x": 2, "score": 0.6},
        {"id": 5, "t": 2, "x": 24, "score": 0.3},
        {"id": 6, "t": 2, "x": 35, "score": 0.7}
]

edges = [
    {"source": 0, "target": 2, "score": 0.9},
    {"source": 1, "target": 3, "score": 0.9},
    {"source": 0, "target": 3, "score": 0.5},
    {"source": 1, "target": 2, "score": 0.5},
    {"source": 2, "target": 4, "score": 0.7},
    {"source": 3, "target": 5, "score": 0.7},
    {"source": 2, "target": 5, "score": 0.3},
    {"source": 3, "target": 4, "score": 0.3},
    {"source": 3, "target": 6, "score": 0.8}
]

graph = nx.DiGraph()
graph.add_nodes_from([
    (cell["id"], cell)
    for cell in cells
])
graph.add_edges_from([
    (edge["source"], edge["target"], edge)
    for edge in edges
])

graph = motile.TrackGraph(graph)

from motile_toolbox.visualization import draw_track_graph

draw_track_graph(graph, alpha_attribute="score", label_attribute="score")

For this tutorial, we will use a simple model: The costs for selecting nodes and edges should depend on their score attribute, and the start of a track should have a positive cost as well. Furthermore, we are only interested in single-object tracks, i.e., tracks don’t merge or split. The solver for this model is:

from motile.constraints import MaxParents, MaxChildren
from motile.costs import NodeSelection, EdgeSelection, Appear

solver = motile.Solver(graph)

solver.add_constraint(MaxParents(1))
solver.add_constraint(MaxChildren(1))

solver.add_cost(NodeSelection(weight=1, attribute="score"))
solver.add_cost(EdgeSelection(weight=-2, attribute="score", constant=1))
solver.add_cost(Appear(constant=1))

Each of those costs is calculated as the product of weights and features:

motile.costs.NodeSelection and motile.costs.EdgeSelection each have a weight (given as argument weight) by which to scale a node or edge feature (given by argument attribute). Both cost terms also support a constant cost (given by argument constant), which is in fact also a weight on an implicit “feature” with value 1.0.

More generally, the cost \(c_y\) of a variable \(y\) is

\[\def\vct#1{\mathbf{#1}} c_y = \vct{w}^\intercal\vct{f}_y \text{.}\]

The variable \(y\) can be an indicator for the selection of a node, the selection of an edge, a node that starts a track, or any other variable added to the solver.

In the example above, the motile.variables.EdgeSelected variable (which is the target of the cost motile.costs.EdgeSelection), has the following weights and features:

\[\begin{split}\vct{w} = \begin{pmatrix} w_\text{attr} \\ w_\text{const} \end{pmatrix} = \begin{pmatrix} -2 \\ 1 \end{pmatrix} \;\;\; \text{and} \;\;\; \vct{f}_e = \begin{pmatrix}\text{attr} \\ 1.0 \end{pmatrix} \text{,}\end{split}\]

where \(\text{attr}\) is the value of the attribute score of edge \(e\).

The motile solver knows about all the weights that have been introduced through cost functions:

solver.weights

('NodeSelection', 'weight') = 1
('NodeSelection', 'constant') = 0.0
('EdgeSelection', 'weight') = -2
('EdgeSelection', 'constant') = 1
('Appear', 'weight') = 1
('Appear', 'constant') = 1

Our initial weights are just a guess, let’s solve…

from motile.variables import NodeSelected, EdgeSelected
from motile_toolbox.visualization import draw_solution

solver.solve()

…and inspect the solution:

print(solver.get_variables(NodeSelected))
print(solver.get_variables(EdgeSelected))

draw_solution(graph, solver, label_attribute="score")

NodeSelected(0): cost=0.800000011920929 value=0.0
NodeSelected(1): cost=0.10000000149011612 value=0.0
NodeSelected(2): cost=0.30000001192092896 value=0.0
NodeSelected(3): cost=0.4000000059604645 value=0.0
NodeSelected(4): cost=0.6000000238418579 value=0.0
NodeSelected(5): cost=0.30000001192092896 value=0.0
NodeSelected(6): cost=0.699999988079071 value=-0.0
EdgeSelected((0, 2)): cost=-0.7999999523162842 value=0.0
EdgeSelected((0, 3)): cost=0.0 value=0.0
EdgeSelected((1, 3)): cost=-0.7999999523162842 value=0.0
EdgeSelected((1, 2)): cost=0.0 value=0.0
EdgeSelected((2, 4)): cost=-0.3999999761581421 value=0.0
EdgeSelected((2, 5)): cost=0.3999999761581421 value=0.0
EdgeSelected((3, 5)): cost=-0.3999999761581421 value=0.0
EdgeSelected((3, 4)): cost=0.3999999761581421 value=0.0
EdgeSelected((3, 6)): cost=-0.6000000238418579 value=-0.0

None of the nodes or edges were selected, which is indeed the cost minimal solution: the cost for selecting nodes or edges is too high.

We can use the solver.weights object to directly modify the weights on the costs. Here, we further lower the cost of edges for example:

solver.weights[("EdgeSelection", "weight")] = -3
solver.weights[("Appear", "constant")] = 0.5

solution = solver.solve()

print(solver.weights)
print(solver.get_variables(NodeSelected))
print(solver.get_variables(EdgeSelected))

draw_solution(graph, solver, label_attribute="score")

('NodeSelection', 'weight') = 1
('NodeSelection', 'constant') = 0.0
('EdgeSelection', 'weight') = -3
('EdgeSelection', 'constant') = 1
('Appear', 'weight') = 1
('Appear', 'constant') = 0.5

NodeSelected(0): cost=0.800000011920929 value=1.0
NodeSelected(1): cost=0.10000000149011612 value=1.0
NodeSelected(2): cost=0.30000001192092896 value=1.0
NodeSelected(3): cost=0.4000000059604645 value=1.0
NodeSelected(4): cost=0.6000000238418579 value=1.0
NodeSelected(5): cost=0.30000001192092896 value=1.0
NodeSelected(6): cost=0.699999988079071 value=0.0
EdgeSelected((0, 2)): cost=-1.6999998092651367 value=1.0
EdgeSelected((0, 3)): cost=-0.5 value=0.0
EdgeSelected((1, 3)): cost=-1.6999998092651367 value=1.0
EdgeSelected((1, 2)): cost=-0.5 value=0.0
EdgeSelected((2, 4)): cost=-1.0999999046325684 value=1.0
EdgeSelected((2, 5)): cost=0.09999996423721313 value=0.0
EdgeSelected((3, 5)): cost=-1.0999999046325684 value=1.0
EdgeSelected((3, 4)): cost=0.09999996423721313 value=0.0
EdgeSelected((3, 6)): cost=-1.4000000953674316 value=0.0

Annotate Ground-Truth

Ground-truth nodes and edges are annotated by giving them a boolean attribute: True for nodes/edges that should be selected, False for nodes/edges that should not be selected, and None (or simply no attribute) for nodes/edges for which it is unknown. The name of the attribute can be freely chosen.

We can do this directly on the track graph. We will add a new attribute gt:

graph.nodes[0]["gt"] = True
graph.nodes[2]["gt"] = True
graph.nodes[4]["gt"] = True
graph.nodes[5]["gt"] = False

graph.edges[(0, 2)]["gt"] = True
graph.edges[(2, 4)]["gt"] = True
graph.edges[(2, 5)]["gt"] = False

The following shows which nodes and edges are part of the ground-truth (green: should be selected, orange: should not be selected):

node_colors = [
  (0, 0, 0) if "gt" not in node else ((0, 128, 0) if node["gt"] else (255, 140, 0))
  for node in graph.nodes.values()
]
edge_colors = [
  (0, 0, 0) if "gt" not in edge else ((0, 128, 0) if edge["gt"] else (255, 140, 0))
  for edge in graph.edges.values()
]

draw_track_graph(
  graph,
  alpha_func=(
    lambda x: "gt" in graph.nodes[x],
    lambda x: "gt" in graph.edges[x]),
  node_color=node_colors,
  edge_color=edge_colors)

Learn Weights

Learning the weights is done by calling motile.Solver.fit_weights() on the ground-truth attribute gt that we just added:

solver.fit_weights(gt_attribute="gt", regularizer_weight=0.01)
optimal_weights = solver.weights

optimal_weights

('NodeSelection', 'weight') = -12.805133780903077
('NodeSelection', 'constant') = -10.883263312243892
('EdgeSelection', 'weight') = -19.908988083473613
('EdgeSelection', 'constant') = -7.255508875000084
('Appear', 'weight') = 0.0
('Appear', 'constant') = 96.37224779784732

To see whether those weights are any good, we will solve again…

solver = motile.Solver(graph)

solver.add_constraint(MaxParents(1))
solver.add_constraint(MaxChildren(1))

solver.add_cost(NodeSelection(weight=1, attribute="score"))
solver.add_cost(EdgeSelection(weight=-2, attribute="score", constant=1))
solver.add_cost(Appear(constant=1))

solver.weights.from_ndarray(optimal_weights.to_ndarray())

solver.solve()

…and look at the solution:

draw_solution(graph, solver, label_attribute="score")

Indeed, the solution from the learnt weights agrees with the ground-truth where given. Note that this is, in general, not the case: for larger graphs and more noisy features than what we use here as a toy example, there might simply be no set of weights that exactly reproduces the ground-truth. In those cases, the learning will at least try to find weights that result in a solution that is as close as possible to the ground-truth.