Repository navigation
Fit the orientations without the strains, on a fused kernel - #19
Merged
Merged
Conversation
…alse), index(fit=True) The refine's steps stop a little short on noisy patterns: on a customer Zenith's Cu scan one neighbour in ten in its single crystal differed by 1.7 degrees. The same Gauss-Newton fit as the strains' but over the three rotations only (7 projections an iteration, not 17) takes the 90th percentile to 0.20 degrees, below the 0.35 of fitting the strains too, at 5,500 patterns a second on a Titan X against 14,400 for the refine alone.
The orientation-only fit projected its 7 candidates into (N, 7, pixels) and reduced them through some 130 torch passes: 34 of a batch's 53 ms on a Titan X. fit_moments projects them and reduces each pattern to 22 raw sums in one pass, the normalized products the Gauss-Newton step needs following in closed form (the derivative of a normalized pattern), and score_ghost scores the line search and the final fit with the ghost term. The 3x3 step is solved by its adjugate, with no batched solver. On the fused path index() takes 1024 patterns a batch (FUSED_BATCH), which stores no per-pixel intermediate. On a customer Zenith's real Cu scan: 14,800 patterns a second against 6,200, the same orientations as the torch fit (median difference 0.000 degrees).
CSSFrancis
changed the title
…block per pattern refine_stepped runs every refine level of a pattern, and fit_orientations every fit iteration, line search, 3x3 solve and the final score, in one block, thread 0 doing the pattern's arithmetic between the block's reductions. They replace fit_moments, score_ghost and the torch arithmetic between them (some 400 small operations a batch). The fit's rotation is the polar part of (I - beta)^-1 in closed form, the turn about w by atan|w|, not an SVD. On a customer Zenith's real Cu scan, frames on the GPU, Titan X: index(fit= True) at 256 patterns a batch 6,100 -> 13,400 a second, at 1024 14,800 -> 16,200; the same orientations as the torch fit (median difference 0.000 degrees).
Steps of a degree and up are scored against the patterns averaged to half the refine shape, a quarter of the pixels: they move the bands by several of its pixels. On a customer Zenith's Cu scan the refine went from 3.9 to 2.5 ms per 256 patterns, index(fit=True) at 1024 a batch from 16,200 to 17,700 a second, with the orientations the fit reached unchanged (median difference 0.000 degrees, none over 5).
fit_orientations samples the master and its gradient once a pixel -- through the bilinear lookup of the master's own column and row derivatives (smooth, not the bilinear facets' slope), the Lambert map and the normalization -- in place of six more samples for central differences. On a customer Zenith's real Cu scan the fit went from 10.2 to 6.1 ms per 256 patterns, index(fit=True) at 1024 a batch from 17,700 to 19,300 a second; its single crystal's neighbours 0.22/0.34 degrees at the 90th/99th percentile against 0.21/0.33, no orientation over 5 degrees from the central-difference fit's. ANALYTIC_JACOBIAN switches back.
fit_passes runs the analytic-Jacobian fit as FIT_SCHEDULE's passes, each scored at the refine's resolution or half it and moved by its own multiple of the Gauss-Newton step, capped at FIT_MAX_STEP. Every pass's sums score where it stands for nothing, so there is no line search and no final pass: the best-scored full-resolution point visited is kept. On a customer Zenith's Cu scan (full, half, half, full) fitted the single crystal's neighbours to 0.19/0.29 degrees at the 90th/99th percentile in 2.5 ms per 256 patterns, where the line-searched fit took 6.1 for 0.22/0.34; index(fit=True) at 1024 a batch 19,300 -> 27,300 a second, at 256 16,300 -> 21,200. An empty FIT_SCHEDULE keeps the line-searched fit.
FIT_SCHEDULE is (full, half, half, half, full): from a refine stopped at 1 degree, on a customer Zenith's Cu scan, its single crystal's neighbours came to 0.17/0.25 degrees at the 90th/99th percentile, the fit 2.5 ms per 256 patterns and index(fit=True) 33,600 a second at 1024 a batch, where a refine on to 0.5 degrees and four passes gave 0.19/0.29 at 28,800.
prepare averages each frame to the refine shape, divides it by the static background, takes off its dynamic background (averaged to a quarter, blurred with reflected edges and drawn back up as the torch path does), averages it to the search shape and normalizes both, in place of some 29 small operations. The same rows to float32 rounding; on a customer Zenith's 512 px Cu frames 2.2 -> 1.1 ms per 256. With the rest, index(fit=True) runs at 37,800 patterns a second at 1024 a batch, 29,300 at 256.
…slope fit_passes is now one template over the first M trace-free numbers: M = 3 is the orientation fit as before, M = 8 (fit_strain_passes) fits the five deviatoric strains with the rotations. Each strain basis matrix is one more fixed direction derivative, so a pass is still one master sample a pixel; the 72 per-pattern sums are reduced by warp shuffles, the 8x8 normal equations are solved by Cholesky on thread 0 in double, and the best point's polar parts are taken on the host as the torch path takes them. The strain fit takes its Jacobian from the bilinear master's own (facet) slope, not the smooth gradient the orientation fit uses: on the real Cu scan the smooth gradient settled about 1e-3 short of the converged torch strain whatever the number of passes, scoring lower, while the facet slope converges to it. Ten full-resolution passes (STRAIN_SCHEDULE) came to 3.3-3.6e-4 of the converged torch fit, against 2.6-3.8e-3 for the torch fit's three iterations, in 9.0-9.2 ms per 256 patterns against 42 ms on a Titan X. index_with_strain takes FUSED_BATCH patterns a batch when the fused fit runs.
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.
Adds an orientation-only fit and makes it fast enough to run on every pattern.
refine_deformation(..., strain=False): the strain fit's Gauss-Newton over the three rotations only;LiveIndexer.index(frames, fit=True)runs it after the stepped refine.refine_steppeddoes every refine level,fit_orientationsevery fit iteration, its line search, the 3x3 solve and the final score, thread 0 doing the pattern's arithmetic between the block's reductions. The fit's normalized products come from 22 raw sums in closed form (the derivative of a normalized pattern), so no(N, 7, pixels)tensor is stored, and its final rotation is the polar part of (I - beta)^-1 in closed form (the turn about w by atan|w|), not an SVD.LiveIndexer.refineusesrefine_steppedon the fused path whether or not a fit follows.prepareruns in one launch, one block per frame: the static and dynamic backgrounds, both shapes and the normalization, the same rows to float32 rounding.FIT_SCHEDULE, default full, half, half, half, full resolution) with no line search and no final pass: every pass's sums score where it stands, and the best-scored full-resolution point is kept. Its step limit isFIT_MAX_STEP. An empty schedule keeps the line-searched fit.ANALYTIC_JACOBIANswitches back to central differences.index(fit=True)takes 1024 patterns a batch (FUSED_BATCH); so doesindex_with_strainwhen the fused strain fit runs (below).Measured on a customer Zenith's real Cu scan (64 x 64, 512 px patterns, 15 kV), frames on the GPU, Titan X, in its single crystal (rows 36 and below) as the misorientation between neighbouring points:
index, stepped refine only (before)index_with_strainindex(fit=True), torch arithmeticindex(fit=True), fusedThe fused rows are from a refine stopped at 1 degree (Autopilot's setting with this fit). The fused fit's orientations are a median 0.2° from the torch fit's on the small grains, none over 5° anywhere in the scan; with central differences and an empty schedule the fused fit is the torch fit to float32 rounding. At 1024 a batch, per 256 patterns (CUDA events): prepare 1.1 ms, search 1.9, refine 1.3, fit 2.6. A 24 px search was tried and left out: 2% faster end to end, but 6.6% of the small grains' points seeded elsewhere and moved more than 5°. A hexagonal phase's dictionary (60,261 entries) costs only its search.
Tests:
test_the_orientation_only_fit_refines_the_orientation_and_no_strain;test_the_fused_prepare_is_the_torch_one,test_the_fused_stepped_refine_is_the_torch_one, andtest_the_fused_orientation_fit_is_the_torch_onefor both Jacobians (CUDA only). At that commit the whole suite passed (397) apart fromtest_distribution_identity, which fails on a stale local install's metadata (0.1.0 against 0.2.1), unrelated.The strain fit on the same kernel
fit_passesis now one CUDA template over the first M trace-free numbers: M = 3 is the orientation fit as before, M = 8 (fit_strain_passes) fits the five deviatoric strains with the three rotations. Each strain basis matrix is one more fixed direction derivative,dv/dp_k = -g B_k (map d), so a pass is still one master sample a pixel. The 72 per-pattern sums (S, S², J[8], sJ[8], the 36 of JJ, Jr[8], Jg[8], s·r, s·g) are reduced by warp shuffles, the 8x8 normal equations (same relative damping as the torch path) are solved by Cholesky on thread 0 in double (a non-positive pivot takes no step), and the kernel returns the best point's eight numbers;polar_partsis then one batched torch op on the host, as the torch path does it. The orientation fit now shares the reduction and the Cholesky solve (its 3x3 adjugate is gone); its real-scan numbers are unchanged (0.17°/0.25°, 30,060/40,570 patterns/s at 256/1024).refine_deformation(..., strain=True)runs it whenindexer.fusedandSTRAIN_SCHEDULEis non-empty (default: ten full-resolution passes, step capSTRAIN_MAX_STEP=FIT_MAX_STEP= 3e-2, 64 threads a block,FIT_STRAIN_THREADS). CPU, MPS and an empty schedule keep the torch path. Like the orientation fit it keeps its best-scored full-resolution point, so only a NaN score is refused (strain NaN, g and score the seed's).Accuracy
Synthetic (
test_the_fused_strain_fit_is_the_torch_one, new; at the indexer's pattern center and with a center map per pattern, a ghost, seeds 0.2° off, 2e-3 strain): fused and torch both recover the strain to 2e-4 and the orientation to 0.02°, and agree with each other to 1e-4 in strain and score. In a 303-orientation sweep at 128 x 128 with 3e-3 strain: noise-free, both 4.5e-5 median error against the truth, fused - torch 7e-7; blurred, both 1.6e-3, fused - torch 1.9e-5; blurred + noise, both 2.1e-3, fused - torch 6.3e-5.test_strain_is_refined_with_the_orientationpasses unchanged (it runs on the CPU, torch path).Real Cu scan, single crystal (rows 36+, recrystallized, 1792 points), the same seeds for every fit. Strain in 1e-4: each component's median, then (IQR, nb) where nb is the median difference between neighbouring points, the map's point-to-point scatter. "converged" is the torch fit at 20 iterations (from the two seed sets below it lands within a few 1e-4 of itself).
Seeds from the full stepped refine (2, 1, 0.5, 0.25, 0.12°):
Seeds from a refine stopped at 1° (Autopilot's setting with the fused fit):
Median |fused - shipped torch| per component (1e-4), full-refine seeds: e11 26.9, e22 10.1, e33 36.2, e12 50.7, e13 27.3, e23 24.6; 1° seeds: 13.9, 40.4, 26.2, 44.0, 70.5, 53.1. Those differences are the three-iteration torch fit not having converged: its strain moves by 26-38e-4 between iterations 3 and 20 and depends on the seed (e13 +34.9 against +78.2), where the fused fit's and the converged fit's do not. Distance from converged by pass count (median, 1° seeds): 6 passes 12.4e-4, 8 passes 5.1e-4, 10 passes 3.6e-4, 12 passes 3.2e-4; the torch fit at 5 iterations 7.6e-4, at 8 iterations about 0.
What this does not show. The strain here is not small: e12 is about 4.4e-3 in every fit, converged or not, which in a recrystallized crystal is the pattern center (and master) error that a single pattern cannot tell from strain, so none of these maps is a validated absolute strain. The spread and the neighbour scatter are the usable precision numbers, and the fused fit's match the converged fit's (and are lower than the shipped fit's in most components).
DEFORMATION_ITERATIONS's docstring records that, on a 15 kV Si indent against CrossCourt, three iterations matched HR-EBSD better than iterating to convergence (rms 3.4-5.4e-4 against 5.0-6.5e-4). The fused fit behaves like the converged one, so if that finding holds generally it would favour the old three-iteration fit; the indent data was not available here and the fused fit has not been checked against an HR-EBSD reference.Speed
index_with_strainon all 4096 frames (on the GPU), Titan X, CUDA events, best of three (range in brackets), and the strain fit alone per 256 patterns. The GPU was shared, butnvidia-smishowed 0-3% utilization around the runs and the repeats agreed to within 3%, so these look uncontaminated.The torch rows are faster than the ~2,100/s quoted for
index_with_strainabove, which was measured differently; the fused/torch ratio is the comparable number: the fit 3.8-4.7x faster,index_with_strain2.5-3.7x. Threads a block for the strain fit, at 1024 a batch: 64 took 9.4 ms per 256, 128 took 10.6, 256 took 13.8.Tests:
test_the_fused_strain_fit_is_the_torch_one(CUDA only, two cases; flipping the sign of one strain derivative in the kernel makes it fail). The whole suite passes (399, 3 skipped) apart fromtest_distribution_identity, as before;ruff checkandruff format --checkpass.Open
refine_deformation'siterations,line_searchandstepdo not apply.