Repository navigation
ENH: Share one IMPACT loss across the metric and ConvexAdam, with masks, 2D models and sampling - #2
Merged
Merged
Conversation
The coarse stage summed the squared differences over every channel of every model, unweighted, while the fine stage takes a weight per layer. A model with larger or more numerous features then drove the coarse search whatever its weight, against a coupling schedule whose coefficients are absolute: beside MIND, the last decoder layer of TotalSegmentator MR (M730) costs 15 times MIND on an abdominal MR/CT pair. SetLayersWeight takes one weight per kept layer, indexed like the fine stage's. Each layer is scaled by the square root of its weight before the cost, which weighs its share of the cost by the weight. A missing entry weighs 1, so existing callers get the same field; a negative weight is refused.
…ration stages Distances are all positive and 0 at a perfect match: Dice, Cosine, L1Cosine and NCC report 1 - similarity. DotProduct, unbounded, is removed. Loss normalization, on by default (NormalizeLosses): each compared layer is divided by its value when a level starts (the metric: the first evaluation after Initialize(); the fine stage: its first iteration; the coarse stage: the cost at zero displacement), so every layer starts at 1 and LayersWeight weighs comparable quantities. The factor lives in Impact::LossNormalization, owned by the caller: bbdd24f removed the earlier normalization because the loss objects, rebuilt at every evaluation, latched it again each time. LNCC: the local normalized cross-correlation of two feature maps, squared as ANTs and FireANTs compute it, over LNCCKernel^Dim windows. The fine stage evaluates it on its dense maps; the point-sampling metric refuses it at Initialize(). The coarse stage takes PCA, and SubsetFeatures is a per-layer number of channels drawn at random in both stages (once in the coarse stage, at every iteration in the fine stage), as the metric draws them. Both stages honour a model's voxel size: the image is resampled to it by linear interpolation before the model sees it; 0 means the image as it is. No smoothing precedes the downsampling: on elastix Static MR/CT a Gaussian of half the voxel size left the Dice unchanged (0.723 against 0.722 over six pairs) and folded more. ImpactModelConfiguration::SetFeatureNormalization (none, l2, standardized) scales each kept layer's feature vectors before PCA, in the stages and in the metric (its Static maps and its online patches).
…eration ImpactFineRegistration::SetSamplingPercentage (in (0, 1], default 1) makes the Static feature similarity read that share of each layer's voxels, drawn anew at every iteration from the seeded generator, as the metric draws its points. Only the drawn points are warped: the fixed features are read at the drawn voxels, the moving ones at the same voxels displaced by the residual field, which is interpolated there as the whole-layer path resizes it (align_corners both), so an iteration costs that share of a full one. Point-wise distances only: LNCC reads windows of whole maps and is refused, as are Jacobian mode and the intensity path, which compare whole images.
A model of lower dimension than the image was swept along the last ITK axis, so a sagittal or coronal acquisition showed the network sagittal or coronal slices where the other engines show it axial ones. Every path that sweeps a volume for a feature map (ImageToFeaturesMap, which serves elastix's Static mode, and the coarse and fine Static stages) now sweeps the image axis most aligned with head-feet, argmax_j |Direction[2][j]|, a tie going to the larger j; an axial image keeps its last axis. RunTiledModel and ExtractFeatureLayers take the axis. A 2D model's PCA was refused. It now follows the FireANTs engine: the basis is fitted on the fixed features of at most 32 slices along the sweep axis, at round(linspace(0, n - 1, min(n, 32))), each image is centred by its own mean over those slices, and every slice is projected (PcaSlices, PcaReduce). A 3D model's PCA is unchanged.
The fine stage's Jacobian mode took a 2D model's slices along z only, chunked by FeatureChunkSize, so the network never saw the anatomy in another orientation. A dense engine draws no points, so it cannot give each point a random plane as the metric does; turning the volume by a random rotation spent 22 to 57 % of the network's work on corners brought in from outside the image (measured in FireANTs). Each iteration now sweeps ONE image axis drawn at random from the seed, chunked along it, and the fixed features of an axis are extracted once and kept, so at most three sets. A 2D model's PCA follows the per-axis rule: the basis is read on the fixed slices along the axis drawn, and the moving features are centred by their own mean over the same slices of the image as the field warps it, taken without the graph as the FireANTs engine does. Each chunk used to be centred by its own mean. A model resampling the image to its voxel size is refused whether or not the volume is chunked: its slices would be compared index by index with slices of another grid, which already failed unchunked. The moving features are no longer extracted at setup in Jacobian mode, which discarded them.
The coarse cost volume was a squared difference per layer whatever the fine stage and the metric compared with. ImpactCoarseRegistration takes SetDistance, one entry per kept layer as in the fine stage (L1, L2, Cosine, L1Cosine, Dice, NCC, LNCC; L2 by default), and each layer's cost at a coarse voxel is its distance over the cell window around it: the two 3^Dim box passes ConvexAdam smooths its cost with. A point-wise distance's terms are averaged over the window, which for L2 is exactly the previous cost; NCC and LNCC both correlate each channel over that window, from its weighted moments of f, m, f^2, m^2 and f*m, since a cost volume has one value per voxel and candidate. The layer weights and the normalization are kept; the normalization latches the same L2 value as before. The per-point terms move into the losses (Loss::forwardPoints, which forwardValue averages) and the windowed correlations into NCC and LNCC, so the three stages share one definition of every distance. The cost volume is built by Impact::AccumulateCoarseCost, which a test checks against loops in double precision for every distance.
SetSamplingPercentage worked in Static mode only; the Jacobian mode ran the network on the whole warped image. Below 100 % the Jacobian mode now follows elastix: at every iteration that share of the fixed voxels is drawn anew from the seed, the points whose every model patch fits in the image are kept (the SampleCheck, read on the patch corners), each model runs on its patch -- its PatchSize, the receptive field, strictly positive on every model axis -- cut in the fixed image and in the moving image as the field warps it, and its layers' centre voxels are compared with the configured distance. A 2D model cuts its patch on a plane drawn for each point, the metric's: PatchPlane's rotation moves into Impact::PatchPlaneRotation, which both use. The gradient reaches the control grid through the warped patches, a batch at a time under the device budget (itkImpactBatchBudget.h), a batch that runs out of memory replayed at half the size with the gradient it had added taken back. A per-point distance weighs each batch by its share of the points; NCC mixes the points, so its gradient with respect to every point's features is taken on the whole set first and seeds each batch's backward. SetBatchSize bounds the batch as the metric's does. PCA follows a 3D model's rule on the points: the basis is fitted on the fixed features of the iteration's points, each side centred by its own mean over them, the moving one without the graph.
The sampled Jacobian mode warped every batch's patches by one smoothed
control grid and back-propagated each batch to the control grid, so the
second batch -- or the second model -- ran through the smoothing's graph
again after the first backward had freed it ("Trying to backward through
the graph a second time", in AvgPool3DBackward0). Tests had used no
smoothing, or one model in one batch.
Every batch now back-propagates into the smoothed field as a detached
leaf, and its accumulated gradient crosses the smoothing passes once, at
the end of the iteration; a batch replayed after running out of device
memory takes back what it had added to that leaf. A test runs two models
on a smoothed half-resolution grid with the normalization on, in several
batches, and checks the first value (each layer at 1) and that the
gradient matches the one of a single batch.
The coarse and fine stages counted their cells, capture range, control grid, step and windows in voxels of each axis, so an anisotropic image got cells, a capture range and a regularisation stretched along its thick axis (on 0.8 x 0.8 x 3 mm: cells of 4.8 x 4.8 x 18 mm, a range of 19 mm in-plane and 72 mm along z, z under-regularised ~4x per mm). Every such knob keeps its number but is read in voxels of the fixed image's finest axis, s_min, and derived per axis, so that it is (nearly) isotropic in millimetres. An isotropic image gives the same numbers and bit-identical fields at any spacing; only anisotropic images change. Coarse stage: a cell spans gs_a = max(1, round(GridSpacing * s_min / s_a)) voxels (per-axis pooling kernel and stride); the capture range is R = DisplacementHalfWidth * GridSpacing * s_min mm, covered along axis a by hw_a = ceil(R / (gs_a * s_a)) cells, candidates prod_a (2 hw_a + 1); the coupling penalty weighs a move of one cell along axis a by (gs_a * s_a / (GridSpacing * s_min))^2. The smoothing and the NCC/LNCC window stay in cells. Fine stage: the control grid is shrink_a = max(1, round(GridShrinkFactor * s_min / s_a)) times coarser along axis a; the displacement is optimized in units of s_min, so an Adam step moves LearningRate * s_min mm along every axis; the regulariser takes the displacement's gradient in mm per mm; the LNCC window is LNCCKernel voxels along the compared map's finest axis in mm and the nearest odd count of the same length along the others (Loss::forwardSpatial takes a side per axis). FeatureChunkSize stays in slices.
Impact::GetPatchIndex, which elastix's IMPACT metric cuts a 2D model's patch with in Jacobian mode, drew each point's plane from the metric's running random generator: the same angle distribution and Rz * Ry * Rx order as the itkv4 metric and ImpactFineRegistration's sampled mode, but not the same planes, and a plane that depended on how many points had been drawn before. It now takes the seed, the model index and the point's fixed-grid index and lays the patch on PatchPlaneRotation's plane, so the three hosts cut identical patches for the same point, seed and model. The running-generator signature is removed; elastix's call site passes the point's index and its seed instead.
Dividing each coarse layer by its cost at zero displacement multiplies a layer with a small cost there into a spread over the candidates far larger than the coupling schedule's absolute coefficients (MR/CT: the TotalSegmentator head by 453, 14x MIND's spread), which is why the normalized coarse stage folded, and why the KonfAI wrapper keeps it raw. SetBalanceLosses weighs the layers on how much each moves the argmin instead: layer l's spread S_l, the mean over the cells of its cost's range over the candidate box, is measured on the first cost volume and the layer is weighed by (sum_k S_k / L) / S_l times its LayersWeight, over the layers kept with a spread above 0. Every layer then spreads alike and their total spread is the raw one, the scale the coupling is calibrated on. It takes the place of NormalizeLosses; off by default. GetLayerSpreads returns the spreads measured.
ImpactCoarseRegistration and ImpactFineRegistration take an optional fixed and moving mask (SetFixedMask, SetMovingMask; Image<unsigned char>, in where not 0, each on its image's grid). Absent, nothing changes: the same code runs, bit for bit, with no mask tensor and no extra operation. Fine stage, as the FireANTs engine: a voxel counts where the fixed mask and the moving mask warped by the current field both hold (>= 0.5), on the grid each distance compares (the masks pooled to it as shares). A point-wise distance averages its terms over those voxels (a ratio of two sums), NCC correlates them, LNCC averages its local terms there, in Static mode, in the intensity path and in both Jacobian modes (the chunked one weighs each chunk by its share of the counted voxels). The sampled modes draw their points in the fixed mask and drop those the warped moving mask does not hold. The moving mask follows the total field when the moving features were refreshed. Coarse stage: a cell's features are its masked voxels' mean (a ratio of two pooled sums), and its cost at a coarse voxel for a candidate is averaged over the cell window weighed by the fixed cell's mask share times the moving mask's share in the cell the candidate shifts it to, moved like the features. A candidate that leaves the window no weight costs as much as the worst candidate with some; a coarse voxel no candidate gives any weight has no data cost, starts at zero displacement and follows its neighbours through the coupling. The normalization and the spread balance read the masked costs.
MeasureBatchBudget read the probe's peak allocation as the memory one patch costs, but that peak includes whatever the device already holds. The metric holds little there; ImpactFineRegistration's sampled Jacobian mode holds the images and the control grid, so half of them counted per patch and the budget fell to a few patches: on an ExaSPIM pair at 80 um (85 M voxels), 0.2 % of the voxels ran in thousands of batches, 24 s an iteration. The probe now counts what it allocates above what was held when it started: 1.9 s an iteration. A CUDA check holds 256 MB while it measures.
SampledPatchPositions formed every patch voxel's position in double precision on the host, for every batch, and copied it to the device. It now takes the tensor options to form them with: the sampled Jacobian mode places its patches (and checks their corners) on its device in single precision, and only a 2D model's per-point planes are still drawn on the host. The positions a test reads stay in double precision.
The budget was the free memory over the cost of one patch, measured at one and two patches. A large network takes more per patch at a few hundred (workspaces, fragmentation): with TotalSegmentator in elastix's metric the budget of 205 patches ran out of memory at the first batch and was halved to 102 for the whole run, slower than before a5cbea4. The estimate is now run once at its size, and scaled down by the measured peak while it takes more than the usable memory (to 3/4 after an out-of-memory), so the batch kept is one that fits (FitBatchToDevice).
Its unbiased deviation over one channel is NaN, which the clamp does not catch, so every distance came out non-finite. KonfAI refuses it alike.
The threader header named itk::Impact::LossNormalization without its header, so it compiled only after another header had included it, which the header test's order did not do.
KWStyle of ITK 5.4 reads 'namespace itk::Impact' as a wrong namespace; the other itkImpact headers nest the two.
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.
One IMPACT loss shared by the ITKv4 metric and the ConvexAdam registration stages, and the capabilities KonfAI's three engines rely on.
GTests: 77 of 78 pass (the CUDA test is skipped on CPU). Needed by KonfAI 1.8.7 as release 0.1.6, with the coarse cost raw by default.