Repository navigation
feat: --xla, XLA-compiled inference for .h5 and .scnm models - #65
Merged
Merged
Conversation
`copick inference easymode --xla` runs each model's forward pass as one XLA-compiled function. It is off by default. An `XLAModel` wrapper is called like the Keras model and passes every other attribute through, so easymode's tiling of an .h5 network and the .scnm inference use it unchanged. Each worker compiles at the first call of each input shape and logs the time. If a compiled call raises, the worker warns once and continues with the plain model. The report records `"xla"`, and each worker lists in `xla_fallbacks` any model that gave up XLA, with the reason. On three 10521 tomograms at TTA 4 on one L40S with TensorFlow 2.21: the microtubule .h5 network went from 310 s to 107 s per tomogram, 2.9x, and the actin .scnm slab model from 87 s to 67 s, 1.3x. The masks matched the plain model's to an IoU of 0.9998 or better. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
easymode 1.2.6 and ais-cryoet 1.2.40 release the --xla options this package mirrors (mgflast/easymode#37, mgflast/Ais#16). The .h5 path uses only _pad_volume, _segment_tomogram_instance and load_model, which are unchanged, so the bump is the supported version and commit, the tests and the docs. The .scnm port follows ais-cryoet 1.2.40: its slab --stride keeps the depth // 2 default, which easymode does not override, and under XLA a short last 2.5D batch is padded with zeros to the batch size so the compiled function keeps one input shape. The docs also drop the claim that easymode is not on PyPI, and name its new metadata pins. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
uermel
pushed a commit
that referenced
this pull request
Oct 21, 2026
🤖 I have created a release *beep* *boop* --- ## [2.0.0-alpha.4](copick-easymode-v2.0.0-alpha.3...copick-easymode-v2.0.0-alpha.4) (2026-10-21) ### Features * --xla, XLA-compiled inference for .h5 and .scnm models ([#65](#65)) ([75ef67b](75ef67b)) --- Merging this PR will create a new alpha prerelease for copick-easymode. Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Port of the main-line
--xlachange and the easymode 1.2.6 bump to v2.0. The code is identical to main's.Adds an opt-in
--xlaflag tocopick inference easymodethat runs each model's forward pass as one XLA-compiled function,tf.function(jit_compile=True). It is the option released upstream aseasymode segment --xlain easymode 1.2.6 (mgflast/easymode#37) andais segment --xlain ais-cryoet 1.2.40 (mgflast/Ais#16).core/xla.py:XLAModelis called like the Keras model and hands every other attribute to the model, so the.h5tiling and the.scnminference use it unchanged. Each input shape compiles once and is logged. If a compiled call raises, the worker warns once and continues with the plain model.InferenceSettings.xla,dispatch_easymode_inference(xla=...)andrun_easymode_inference(xla=...)carry it to the workers.--xla/--no-xla, off by default, because a missing or too oldptxascan abort TensorFlow outright."xla", and per workerxla_fallbackswithfeatureandreason.v1.2.6(00ab903).tagged-easymodeCI job andVALIDATION.mdinstall it..h5functions this package uses are unchanged..scnmport follows ais-cryoet 1.2.40.Measured on the main line: three 10521 tomograms at TTA 4 on one L40S with TensorFlow 2.21.
.h5went from 310 s to 107 s per tomogram, 2.9x. Actin.scnmslab went from 87 s to 67 s, 1.3x.Tests:
00ab903.🤖 Generated with Claude Code