Seeded watershed performance

[PerretB]

Improve efficiency of labelsitation_seeded_watersed

Higra/include/higra/algo/watershed.hpp

Line 124 in f85538d

auto labelisation_seeded_watershed(

Idea 1: sure, shortcut first loop when num_vertices - num_seeds edges have been processed (edge case, num_seeds = 0)
Idea 2: to test, replace relabelling of nonlabeled vertices with union-find by a connected component labelling

[lapertor]

Idea 1 has been explored and implemented in Pull Request #296.

Today I explored the connected-component relabelling idea (Idea 2).

I implemented a variant that propagates labels incrementally during union-find merges, instead of doing the final relabelling pass at the end. To do that, I maintained explicit component member lists so that vertices of a background component could be updated when merging with a labelled component.

Current status:

• passes the existing test_watershed.py tests
• passes a few additional small correctness checks I wrote locally
• but introduces a performance regression on my benchmarks

For a 2000x2000 4-adjacency graph:

• current implementation: watershed loop around 0.85s on most cases, and around 0.69s when all vertices are seeded
• experimental implementation: watershed loop around 2.4s on most cases, and around 1.74s when all vertices are seeded

So this variant seems functionally correct, but slower than the current implementation, likely because maintaining component member lists and propagating labels during merges is more expensive than the final linear relabelling pass.

Because of that, I do not think this is a good PR candidate in its current form.

Benchmark script: see #296.
Old implementation (master, not Idea 1): see #296.
Output of tests on the old implementation (master, not Idea 1): see #296.

New implementation (variant of Idea 2)

    template<typename graph_t, typename T1, typename T2>
    auto labelisation_seeded_watershed(
            const graph_t &graph,
            const xt::xexpression<T1> &xedge_weights,
            const xt::xexpression<T2> &xvertex_seeds,
            const typename T2::value_type background_label = 0) {
        HG_TRACE();
        auto &edge_weights = xedge_weights.derived_cast();
        auto &vertex_seeds = xvertex_seeds.derived_cast();
        hg_assert_edge_weights(graph, edge_weights);
        hg_assert_node_weights(graph, vertex_seeds);
        hg_assert_1d_array(edge_weights);
        hg_assert_1d_array(vertex_seeds);

        using label_type = typename T2::value_type;

        array_1d<index_t> sorted_edges_indices = stable_arg_sort(edge_weights);

        auto t0 = std::chrono::high_resolution_clock::now();

        index_t num_nodes = num_vertices(graph);
        index_t num_edges = sorted_edges_indices.size();

        union_find uf(num_nodes);

        array_1d<label_type> labels = vertex_seeds;
        array_1d<label_type> component_labels = vertex_seeds;
        array_1d<index_t> component_head = xt::arange<index_t>(num_nodes);
        array_1d<index_t> component_tail = xt::arange<index_t>(num_nodes);
        array_1d<index_t> component_next = array_1d<index_t>::from_shape({static_cast<size_t>(num_nodes)});

        index_t invalid = num_nodes;
        std::fill(component_next.begin(), component_next.end(), invalid);

        for (index_t i = 0; i < num_edges; i++) {
            auto ei = sorted_edges_indices[i];
            auto e = edge_from_index(ei, graph);
            auto c1 = uf.find(source(e, graph));
            auto c2 = uf.find(target(e, graph));

            if (c1 == c2) {
                continue;
            }

            auto l1 = component_labels(c1);
            auto l2 = component_labels(c2);

            if (l1 != background_label && l2 != background_label && l1 != l2) {
                continue;
            }

            label_type merged_label = l1;
            if (merged_label == background_label) {
                merged_label = l2;
            }

            index_t background_component = invalid;
            if (l1 == background_label && l2 != background_label) {
                background_component = c1;
            } else if (l2 == background_label && l1 != background_label) {
                background_component = c2;
            }

            if (background_component != invalid) {
                auto current = component_head(background_component);
                while (current != invalid) {
                    labels(current) = merged_label;
                    current = component_next(current);
                }
            }

            auto new_root = uf.link(c1, c2);
            auto old_root = new_root == c1 ? c2 : c1;

            auto new_head = component_head(new_root);
            auto new_tail = component_tail(new_root);
            auto old_head = component_head(old_root);
            auto old_tail = component_tail(old_root);

            component_labels(new_root) = merged_label;
            component_next(new_tail) = old_head;
            component_tail(new_root) = old_tail;

            component_head(old_root) = invalid;
            component_tail(old_root) = invalid;
        }

        auto t1 = std::chrono::high_resolution_clock::now();
        std::cerr << "watershed loop: "
                << std::chrono::duration<double, std::milli>(t1 - t0).count()
                << " ms" << std::endl;

        return labels;
    };

Output of tests on the new implementation

============================================================
Shape: (2000, 2000)
Block shape: (2000, 2000)
Num vertices: 4000000
Num edges: 7996000
Num seeds: 1
watershed loop: 2382.5 ms
watershed loop: 2396.99 ms
watershed loop: 2404.8 ms
watershed loop: 2406.88 ms
watershed loop: 2396.18 ms
Times: [7.604, 7.5841, 7.629, 7.6459, 7.6168]
Mean: 7.6160 s
Std:  0.0211 s
Labels range: 1 to 1

============================================================
Shape: (2000, 2000)
Block shape: (100, 100)
Num vertices: 4000000
Num edges: 7996000
Num seeds: 400
watershed loop: 2379.5 ms
watershed loop: 2421.71 ms
watershed loop: 2412.67 ms
watershed loop: 2390.5 ms
watershed loop: 2388.8 ms
Times: [12.7183, 12.749, 12.8848, 12.8799, 12.6951]
Mean: 12.7854 s
Std:  0.0810 s
Labels range: 1 to 400

============================================================
Shape: (2000, 2000)
Block shape: (10, 10)
Num vertices: 4000000
Num edges: 7996000
Num seeds: 40000
watershed loop: 2400.48 ms
watershed loop: 2421.01 ms
watershed loop: 2399.98 ms
watershed loop: 2391.75 ms
watershed loop: 2401.54 ms
Times: [13.6586, 13.7104, 13.699, 13.5339, 13.532]
Mean: 13.6268 s
Std:  0.0785 s
Labels range: 1 to 40000

============================================================
Shape: (2000, 2000)
Block shape: (4, 4)
Num vertices: 4000000
Num edges: 7996000
Num seeds: 250000
watershed loop: 2390.27 ms
watershed loop: 2388.84 ms
watershed loop: 2385.14 ms
watershed loop: 2407.41 ms
watershed loop: 2399.74 ms
Times: [13.6554, 13.5871, 13.5451, 13.6757, 13.4227]
Mean: 13.5772 s
Std:  0.0904 s
Labels range: 1 to 250000

============================================================
Shape: (2000, 2000)
Block shape: (2, 2)
Num vertices: 4000000
Num edges: 7996000
Num seeds: 1000000
watershed loop: 2311.67 ms
watershed loop: 2303.61 ms
watershed loop: 2385.88 ms
watershed loop: 2319.23 ms
watershed loop: 2300.54 ms
Times: [12.5828, 12.5449, 12.7635, 12.8515, 12.8602]
Mean: 12.7206 s
Std:  0.1329 s
Labels range: 1 to 1000000

============================================================
Shape: (2000, 2000)
Block shape: (1, 1)
Num vertices: 4000000
Num edges: 7996000
Num seeds: 4000000
watershed loop: 1738.24 ms
watershed loop: 1731.6 ms
watershed loop: 1739.08 ms
watershed loop: 1736.71 ms
watershed loop: 1735.14 ms
Times: [6.9304, 6.9384, 6.9941, 6.9738, 6.9852]
Mean: 6.9644 s
Std:  0.0254 s
Labels range: 1 to 4000000

[lapertor]

I will explore Idea 2 (not a variant this time) too later.

[PerretB]

Maybe you should first measure the time spent in the final relabelling loops versus the time spent in the main union-find loop and sorting. The second idea in my original post might be irrelevant if the time of the relabelling is small versus the rest of it.

Note that the time for sorting might vary greatly depending of your compile option : with TBB sorting will be parallelized, which provides large speedups (and of course don't forget to use Release build with cmake)

[lapertor]

You're right, I should have begun with this first.

I did a quick instrumentation pass on master to measure the time spent in the main parts of labelisation_seeded_watershed.

On my 2000x2000 4-adjacency benchmarks, the timings are consistently dominated by the sorting step, then by the main union-find loop, while the final relabelling pass is comparatively small.

Typical numbers are:

• sort: about 5 to 11.8 s
• union-find loop: about 0.53 to 0.70 s
• final relabelling: about 0.06 to 0.15 s

For example, on the (100, 100) block case (0.01 % seeds):

• sort = 10.6 to 10.9 s
• union = 0.66 to 0.69 s
• final relabel = 0.14 s

So the final relabelling looks negligible compared to the sorting, and still small compared to the union-find, I think we can stop exploring that direction then.

Instrumentation code

    template<typename graph_t, typename T1, typename T2>
    auto labelisation_seeded_watershed(
            const graph_t &graph,
            const xt::xexpression<T1> &xedge_weights,
            const xt::xexpression<T2> &xvertex_seeds,
            const typename T2::value_type background_label = 0) {

        auto t0 = std::chrono::high_resolution_clock::now();

        HG_TRACE();
        auto &edge_weights = xedge_weights.derived_cast();
        auto &vertex_seeds = xvertex_seeds.derived_cast();
        hg_assert_edge_weights(graph, edge_weights);
        hg_assert_node_weights(graph, vertex_seeds);
        hg_assert_1d_array(edge_weights);
        hg_assert_1d_array(vertex_seeds);

        using label_type = typename T2::value_type;

        auto t1 = std::chrono::high_resolution_clock::now();

        array_1d<index_t> sorted_edges_indices = stable_arg_sort(edge_weights);

        auto t2 = std::chrono::high_resolution_clock::now();

        index_t num_nodes = num_vertices(graph);
        index_t num_edges = sorted_edges_indices.size();

        union_find uf(num_nodes);

        array_1d<label_type> labels = vertex_seeds;

        auto t3 = std::chrono::high_resolution_clock::now();

        for (index_t i = 0; i < num_edges; i++) {
            auto ei = sorted_edges_indices[i];
            auto e = edge_from_index(ei, graph);
            auto c1 = uf.find(source(e, graph));
            auto c2 = uf.find(target(e, graph));

            if (c1 != c2 && (labels(c1) == background_label || labels(c2) == background_label)) {
                if (labels(c1) == background_label) {
                    labels(c1) = labels(c2);
                } else {
                    labels(c2) = labels(c1);
                }
                uf.link(c1, c2);
            }

        }

        auto t4 = std::chrono::high_resolution_clock::now();

        for (index_t i = 0; i < num_nodes; i++) {
            if (labels(i) == background_label) {
                labels(i) = labels(uf.find(i));
            }
        }

        auto t5 = std::chrono::high_resolution_clock::now();

        auto sort_ms = std::chrono::duration_cast<std::chrono::milliseconds>(t2 - t1).count();
        auto union_ms = std::chrono::duration_cast<std::chrono::milliseconds>(t4 - t3).count();
        auto relabel_ms = std::chrono::duration_cast<std::chrono::milliseconds>(t5 - t4).count();
        auto total_ms = std::chrono::duration_cast<std::chrono::milliseconds>(t5 - t0).count();

        std::cerr << "[seeded_watershed] sort=" << sort_ms
                  << "ms union=" << union_ms
                  << "ms relabel=" << relabel_ms
                  << "ms total=" << total_ms
                  << "ms" << std::endl;

        return labels;
    };

Output

============================================================
Shape: (2000, 2000)
Block shape: (2000, 2000)
Num vertices: 4000000
Num edges: 7996000
Num seeds: 1
[seeded_watershed] sort=5087ms union=649ms relabel=140ms total=5950ms
[seeded_watershed] sort=5096ms union=659ms relabel=140ms total=5970ms
[seeded_watershed] sort=5081ms union=653ms relabel=141ms total=5948ms
[seeded_watershed] sort=5181ms union=661ms relabel=144ms total=6060ms
[seeded_watershed] sort=5181ms union=654ms relabel=142ms total=6051ms
Times: [5.9664, 5.9894, 5.9699, 6.078, 6.072]
Mean: 6.0151 s
Std:  0.0495 s
Labels range: 1 to 1

============================================================
Shape: (2000, 2000)
Block shape: (100, 100)
Num vertices: 4000000
Num edges: 7996000
Num seeds: 400
[seeded_watershed] sort=10567ms union=674ms relabel=147ms total=11464ms
[seeded_watershed] sort=10632ms union=663ms relabel=139ms total=11507ms
[seeded_watershed] sort=10633ms union=675ms relabel=143ms total=11526ms
[seeded_watershed] sort=10900ms union=687ms relabel=144ms total=11810ms
[seeded_watershed] sort=10965ms union=687ms relabel=142ms total=11871ms
Times: [11.4797, 11.5256, 11.5449, 11.8289, 11.8912]
Mean: 11.6541 s
Std:  0.1706 s
Labels range: 1 to 400

============================================================
Shape: (2000, 2000)
Block shape: (10, 10)
Num vertices: 4000000
Num edges: 7996000
Num seeds: 40000
[seeded_watershed] sort=11310ms union=665ms relabel=142ms total=12191ms
[seeded_watershed] sort=11317ms union=691ms relabel=143ms total=12232ms
[seeded_watershed] sort=11725ms union=702ms relabel=148ms total=12651ms
[seeded_watershed] sort=11791ms union=693ms relabel=151ms total=12712ms
[seeded_watershed] sort=11656ms union=682ms relabel=145ms total=12557ms
Times: [12.2077, 12.2521, 12.6718, 12.7299, 12.5783]
Mean: 12.4880 s
Std:  0.2166 s
Labels range: 1 to 40000

============================================================
Shape: (2000, 2000)
Block shape: (4, 4)
Num vertices: 4000000
Num edges: 7996000
Num seeds: 250000
[seeded_watershed] sort=11378ms union=695ms relabel=149ms total=12299ms
[seeded_watershed] sort=11322ms union=688ms relabel=148ms total=12234ms
[seeded_watershed] sort=11326ms union=693ms relabel=148ms total=12243ms
[seeded_watershed] sort=11304ms union=685ms relabel=148ms total=12212ms
[seeded_watershed] sort=11182ms union=679ms relabel=148ms total=12084ms
Times: [12.3148, 12.2517, 12.2618, 12.2315, 12.103]
Mean: 12.2325 s
Std:  0.0704 s
Labels range: 1 to 250000

============================================================
Shape: (2000, 2000)
Block shape: (2, 2)
Num vertices: 4000000
Num edges: 7996000
Num seeds: 1000000
[seeded_watershed] sort=10576ms union=691ms relabel=148ms total=11493ms
[seeded_watershed] sort=10651ms union=704ms relabel=145ms total=11577ms
[seeded_watershed] sort=10467ms union=665ms relabel=153ms total=11360ms
[seeded_watershed] sort=15256ms union=1198ms relabel=247ms total=16856ms
[seeded_watershed] sort=14265ms union=665ms relabel=140ms total=15145ms
Times: [11.5097, 11.5971, 11.3792, 16.8945, 15.174]
Mean: 13.3109 s
Std:  2.2903 s
Labels range: 1 to 1000000

============================================================
Shape: (2000, 2000)
Block shape: (1, 1)
Num vertices: 4000000
Num edges: 7996000
Num seeds: 4000000
[seeded_watershed] sort=5211ms union=551ms relabel=66ms total=5905ms
[seeded_watershed] sort=5230ms union=549ms relabel=65ms total=5920ms
[seeded_watershed] sort=5166ms union=542ms relabel=65ms total=5849ms
[seeded_watershed] sort=5079ms union=531ms relabel=62ms total=5747ms
[seeded_watershed] sort=5066ms union=535ms relabel=64ms total=5742ms
Times: [5.923, 5.9401, 5.8666, 5.7659, 5.76]
Mean: 5.8511 s
Std:  0.0760 s
Labels range: 1 to 4000000

[PerretB]

I'm a bit surprised by the values you're getting for sorting. I did a quick test on my machine ( i9-13900HX, VS 2026, c++20) with stable_sort on a random array of 2 * 2000 * 2000 elements (roughly the number of edges in your tests), and I'm getting the following results:

• Debug: ~2s
• Release: ~0.5s
• Release + parallel execution : ~0.1s

What are the specs of your computer, and what build configuration are you using?

[lapertor]

I checked my setup and you're right, I had actually benchmarked a Debug build.

My setup is:

• CPU: Intel Core i7-14700KF
• 20 cores / 28 logical processors
• 32 GB RAM
• MSVC 19.50.35727 for x64
• Windows, Ninja generator, cl.exe

From CMakeCache.txt, my previous build was CMAKE_BUILD_TYPE=Debug.

I reran everything in a Release build, and the numbers are now more in line with what you reported. On the same 2000x2000 4-adjacency benchmarks, I now get roughly:

• sort: ~230 to 355 ms
• union: ~12 to 25 ms
• relabel: ~1 to 4 ms
• total: ~267 to 400 ms

So the previous 12 s timings were probably due to benchmarking a Debug build. Also, the conclusion remains the same in Release: the final relabelling step is very small compared to the sorting step.

Intrumentation code: same as before

Output (in Release)

============================================================
Shape: (2000, 2000)
Block shape: (2000, 2000)
Num vertices: 4000000
Num edges: 7996000
Num seeds: 1
[seeded_watershed] sort=230ms union=15ms relabel=2ms total=267ms
[seeded_watershed] sort=228ms union=15ms relabel=3ms total=269ms
[seeded_watershed] sort=231ms union=14ms relabel=2ms total=267ms
[seeded_watershed] sort=232ms union=16ms relabel=2ms total=271ms
[seeded_watershed] sort=229ms union=15ms relabel=2ms total=267ms
Times: [0.2765, 0.2788, 0.2791, 0.2831, 0.2757]
Mean: 0.2786 s
Std:  0.0026 s
Labels range: 1 to 1

============================================================
Shape: (2000, 2000)
Block shape: (100, 100)
Num vertices: 4000000
Num edges: 7996000
Num seeds: 400
[seeded_watershed] sort=315ms union=15ms relabel=2ms total=354ms
[seeded_watershed] sort=316ms union=16ms relabel=2ms total=355ms
[seeded_watershed] sort=320ms union=16ms relabel=2ms total=360ms
[seeded_watershed] sort=320ms union=15ms relabel=2ms total=358ms
[seeded_watershed] sort=319ms union=15ms relabel=2ms total=358ms
Times: [0.3629, 0.3673, 0.3711, 0.3676, 0.3686]
Mean: 0.3675 s
Std:  0.0027 s
Labels range: 1 to 400

============================================================
Shape: (2000, 2000)
Block shape: (10, 10)
Num vertices: 4000000
Num edges: 7996000
Num seeds: 40000
[seeded_watershed] sort=340ms union=23ms relabel=2ms total=387ms
[seeded_watershed] sort=344ms union=21ms relabel=2ms total=389ms
[seeded_watershed] sort=344ms union=20ms relabel=2ms total=388ms
[seeded_watershed] sort=341ms union=21ms relabel=2ms total=385ms
[seeded_watershed] sort=341ms union=21ms relabel=2ms total=387ms
Times: [0.3972, 0.3997, 0.3996, 0.3957, 0.3969]
Mean: 0.3978 s
Std:  0.0016 s
Labels range: 1 to 40000

============================================================
Shape: (2000, 2000)
Block shape: (4, 4)
Num vertices: 4000000
Num edges: 7996000
Num seeds: 250000
[seeded_watershed] sort=355ms union=21ms relabel=2ms total=400ms
[seeded_watershed] sort=341ms union=21ms relabel=2ms total=386ms
[seeded_watershed] sort=345ms union=21ms relabel=4ms total=390ms
[seeded_watershed] sort=338ms union=21ms relabel=2ms total=383ms
[seeded_watershed] sort=337ms union=25ms relabel=3ms total=385ms
Times: [0.4086, 0.3953, 0.4025, 0.3956, 0.3949]
Mean: 0.3994 s
Std:  0.0054 s
Labels range: 1 to 250000

============================================================
Shape: (2000, 2000)
Block shape: (2, 2)
Num vertices: 4000000
Num edges: 7996000
Num seeds: 1000000
[seeded_watershed] sort=324ms union=23ms relabel=2ms total=371ms
[seeded_watershed] sort=321ms union=22ms relabel=2ms total=366ms
[seeded_watershed] sort=329ms union=24ms relabel=2ms total=378ms
[seeded_watershed] sort=321ms union=23ms relabel=2ms total=367ms
[seeded_watershed] sort=330ms union=23ms relabel=3ms total=378ms
Times: [0.3794, 0.3775, 0.3872, 0.3789, 0.3919]
Mean: 0.3830 s
Std:  0.0056 s
Labels range: 1 to 1000000

============================================================
Shape: (2000, 2000)
Block shape: (1, 1)
Num vertices: 4000000
Num edges: 7996000
Num seeds: 4000000
[seeded_watershed] sort=232ms union=14ms relabel=1ms total=268ms
[seeded_watershed] sort=230ms union=15ms relabel=2ms total=268ms
[seeded_watershed] sort=236ms union=12ms relabel=1ms total=271ms
[seeded_watershed] sort=236ms union=13ms relabel=1ms total=273ms
[seeded_watershed] sort=242ms union=13ms relabel=1ms total=278ms
Times: [0.2769, 0.2777, 0.2814, 0.2824, 0.2877]
Mean: 0.2812 s
Std:  0.0039 s
Labels range: 1 to 4000000