Repository navigation
Conversation
The v4 metric drew its feature subset (SubsetFeatures) from a generator per work unit, reseeded at every evaluation: per point in Static mode, per batch in Jacobian mode. A point's channels therefore depended on its rank in its work unit, on the batch size, and on whether a batch was replayed after running out of device memory. NCC accumulates its statistics channel by channel over the points, so its sums mixed channels. With 2 of 4 channels, the Jacobian derivative moved by 7 to 15% between batch sizes. The subsets are now drawn once per evaluation, before the threads start, from the seed the planes use, and every point compares the same channels. Two evaluations of the same parameters still compare the same channels. The batch-size test now keeps 2 of the 4 channels instead of all of them, and a new test compares 1 and 4 work units in Static mode. Both fail without this change.
Evaluate and EvaluateDerivative took both by value, so every call copied the subset, two or three times per point and per layer in Static mode.
Both stages called torch::manual_seed(Seed), which reseeds every torch generator of the process: a Python caller's random numbers restarted from the stage's seed after each run. The fine stage also drew its sampled points on the device, from the device's global generator, which the batch-budget probe draws from too, so a probe could shift them. Each stage now draws from a CPU generator of its own, seeded with Seed, and moves the draws to the device. A seed gives the same channels and points on the CPU and on a GPU. The new test checks that running either stage leaves torch's global generator where it was; it fails without this change.
The point-wise losses (L1, L2, Cosine, L1Cosine, Dice) summed a batch's values in float32 before adding the total to their double accumulator, so the value lost digits with the size of the batch: 2e-8 to 6e-8 of relative error over 20000 points. The online mode of the v4 metric and elastix's IMPACT metric both hand the losses whole batches, and finite differences of such a value are mostly rounding. The batch is now summed in double. The new test compares a batch of 20000 points with the same points fed one by one, and fails without this change.
In Static mode the v4 threader let ITK walk the points one by one. Each point went through ITK's ProcessVirtualPoint, which also computed an image gradient IMPACT never reads, then through ProcessPoint, which built a handful of one-row tensors to update the losses. Static now gathers a work unit's points, as the online mode already did, and evaluates them in batches: the feature vectors and the transform Jacobians are written into contiguous tensors, and each loss is updated once per batch. A batch holds at most about 4 MB of [points, channels, parameters] Jacobian, so that a dense domain under a B-spline transform stays bounded. The vector interpolator gains the buffer variants this needs, and a value-and-derivative evaluation that shares the B-spline weights between the two. On the lung CT test data, with 20000 points and 8 of a model's 32 channels, an evaluation takes 11 ms instead of 40 ms on 24 threads, and 136 ms instead of 468 ms on one.
ImageToImageMetricv4 smooths the whole moving image into a gradient image in Initialize(), for the metrics that read image gradients at their points. IMPACT compares features: the online mode takes its patch gradients from the moving interpolator, and Static mode differentiates the feature maps. The filter is now off by default. On the lung CT, Initialize() of an online metric takes 0.01 s instead of 0.08 s.
InterpolateVectorImageFunction held one BSplineInterpolateImageFunction per channel. Every channel of every point converted the point to a continuous index, computed the support, the weights and the mirrored indices, and allocated two or three vnl matrices to hold them, although all of these are the same for every channel. It now computes each channel's coefficients as that class does, and evaluates all the selected channels with one support and one set of weights, in the same order of operations, so the result is the same to the bit. SetActiveChannels() keeps a copy of the channels an evaluation compares, the channels of a voxel side by side, so that a point reads the fewest cache lines whether it compares all of a model's channels or a few; the v4 threader sets it once per evaluation. The per-channel images are gone, and a second SetInputImage() no longer keeps the first image's channels. The interpolator test now compares values and derivatives with BSplineInterpolateImageFunction channel by channel: 70 channels, a region that does not start at zero, anisotropic spacing, a rotated direction, through the active channels and through another subset. On 32 channels, a value takes 1.2 us instead of 8.4 us, a value and its derivative 2.5 us instead of 19.5 us.
On a GPU the online mode ran each work unit's points as batches of at most BatchSize (64 by default), one work unit after the other under the device lock. The model takes 20000 patches in about 0.08 s, but they went through it as some 300 batches. On a GPU the work units now only collect their points. Once they are done, all the threads cut the patches of a batch, and the models run on batches as large as the device holds. BatchSize now defaults to 0, which means that measured bound on a GPU, and 64 on the CPU, where the work units remain the parallelism and keep evaluating their own points. The measurement tries no batch larger than the largest an evaluation can run: a trial batch sized to the whole free memory took it from the other processes on a shared GPU. On the lung CT test data, with the 32-channel model, an evaluation of 20000 points takes 112 ms instead of 430 ms on a GPU; with 2000 points, 13 ms instead of 103 ms (L2) and 94 ms instead of 782 ms (NCC). The new test compares the GPU's value and derivative with the CPU's, and is skipped without a CUDA device.
NCC compares statistics taken over all the points, so its derivative with respect to one point's features is only known once every point has been seen. The online mode took every channel's d(feature)/d(moving coordinate) instead, from one backward pass per channel: 32 passes for a 32-channel layer. A global loss now takes two passes. The first, without gradient, computes the features of every point, from which the loss takes its value and its derivative with respect to each point's features. The second seeds one backward pass per batch with that derivative, as for L2. To see every point at once, the work units now only collect their points on the CPU too; the threads then evaluate them in shares, 64 points at a time as before. On the lung CT test data, with 2000 points and the 32-channel model, an NCC evaluation takes 15 ms instead of 94 ms on a GPU, and 121 ms instead of 1006 ms on the CPU. L2 is unchanged.
vboussot
changed the title
This branch has not been deployed
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
1 participant
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.
Fixes how the v4 metric and the ConvexAdam stages draw their random numbers, and makes the v4 metric faster: Static mode about 6x, the online (Jacobian) mode 4x to 50x on a GPU.
Fixes
torch::manual_seed(Seed), which reseeds the whole process, a Python caller's generator included. They now use a generator of their own.Speed, Static mode (lung CT test data, 20000 points, a 32-channel model)
BSplineInterpolateImageFunction, which the new interpolator test checks.Initialize().Speed, online mode (same data and model, 2000 points unless noted, best of 3)
BatchSizenow defaults to 0: that measured bound on a GPU, 64 on the CPU.Each fix comes with a test that fails without it. A new test checks that the GPU gives the CPU's value and derivative, for L2 and NCC; it is skipped without a CUDA device.