diff --git a/include/ImpactLoss.h b/include/ImpactLoss.h index 817b684..5e3010e 100644 --- a/include/ImpactLoss.h +++ b/include/ImpactLoss.h @@ -314,7 +314,7 @@ class L1 : public Loss updateValue(torch::Tensor & fixedOutput, torch::Tensor & movingOutput) override { this->initialize(fixedOutput); - this->m_value += (fixedOutput - movingOutput).abs().mean(1).sum().item(); + this->m_value += (fixedOutput - movingOutput).abs().mean(1).sum(torch::kDouble).item(); } torch::Tensor @@ -322,7 +322,7 @@ class L1 : public Loss { this->initialize(fixedOutput); torch::Tensor diffOutput = fixedOutput - movingOutput; - this->m_value += diffOutput.abs().mean(1).sum().item(); + this->m_value += diffOutput.abs().mean(1).sum(torch::kDouble).item(); return -torch::sign(diffOutput) / fixedOutput.size(1); } @@ -359,7 +359,7 @@ class L2 : public Loss updateValue(torch::Tensor & fixedOutput, torch::Tensor & movingOutput) override { this->initialize(fixedOutput); - this->m_value += (fixedOutput - movingOutput).pow(2).mean(1).sum().item(); + this->m_value += (fixedOutput - movingOutput).pow(2).mean(1).sum(torch::kDouble).item(); } torch::Tensor @@ -367,7 +367,7 @@ class L2 : public Loss { this->initialize(fixedOutput); torch::Tensor diffOutput = fixedOutput - movingOutput; - this->m_value += diffOutput.pow(2).mean(1).sum().item(); + this->m_value += diffOutput.pow(2).mean(1).sum(torch::kDouble).item(); return -2 * diffOutput / fixedOutput.size(1); } @@ -420,7 +420,7 @@ class Dice : public Loss torch::Tensor dice = 2.0 * intersectionSum / unionSumSafe; dice.masked_fill_(isEmpty, 1.0); // empty/empty => Dice = 1 - this->m_value += (1.0 - dice).sum().item(); + this->m_value += (1.0 - dice).sum(torch::kDouble).item(); } torch::Tensor @@ -435,7 +435,7 @@ class Dice : public Loss torch::Tensor dice = 2.0 * intersectionSum / unionSumSafe; dice.masked_fill_(isEmpty, 1.0); - this->m_value += (1.0 - dice).sum().item(); + this->m_value += (1.0 - dice).sum(torch::kDouble).item(); torch::Tensor grad = -2.0 * (fixedOutput * unionSumSafe.unsqueeze(-1) - intersectionSum.unsqueeze(-1)) / (unionSumSafe * unionSumSafe).unsqueeze(-1); @@ -492,7 +492,7 @@ class L1Cosine : public Loss torch::Tensor norm_moving = torch::norm(movingOutput, 2, 1); torch::Tensor cosine = dot_product / (norm_fixed * norm_moving); torch::Tensor expL1 = torch::exp(-this->lambda * (fixedOutput - movingOutput).abs()); - this->m_value += (1.0 - (cosine.unsqueeze(-1) * expL1).mean(1)).sum().item(); + this->m_value += (1.0 - (cosine.unsqueeze(-1) * expL1).mean(1)).sum(torch::kDouble).item(); } torch::Tensor @@ -515,7 +515,7 @@ class L1Cosine : public Loss (dot_product.unsqueeze(-1) * movingOutput) / (v * norm_moving.pow(2)).unsqueeze(-1)); torch::Tensor expSum = expL1.sum(1).unsqueeze(-1); torch::Tensor dExpL1 = this->lambda * torch::sign(diffOutput) * expL1; - this->m_value += (1.0 - (cosine.unsqueeze(-1) * expL1).mean(1)).sum().item(); + this->m_value += (1.0 - (cosine.unsqueeze(-1) * expL1).mean(1)).sum(torch::kDouble).item(); return (minusDcos * expSum - cosine.unsqueeze(-1) * dExpL1) / fixedOutput.size(1); } @@ -558,7 +558,7 @@ class Cosine : public Loss torch::Tensor dot_product = (fixedOutput * movingOutput).sum(1); torch::Tensor norm_fixed = torch::norm(fixedOutput, 2, 1); torch::Tensor norm_moving = torch::norm(movingOutput, 2, 1); - this->m_value += (1.0 - dot_product / (norm_fixed * norm_moving)).sum().item(); + this->m_value += (1.0 - dot_product / (norm_fixed * norm_moving)).sum(torch::kDouble).item(); } torch::Tensor @@ -569,7 +569,7 @@ class Cosine : public Loss torch::Tensor norm_fixed = torch::norm(fixedOutput, 2, 1); torch::Tensor norm_moving = torch::norm(movingOutput, 2, 1); torch::Tensor v = (norm_fixed * norm_moving); - this->m_value += (1.0 - dot_product / v).sum().item(); + this->m_value += (1.0 - dot_product / v).sum(torch::kDouble).item(); // d(-cosine)/d(movingOutput_c) = -( f_c/v - (f.m) m_c / (v |m|^2) ). The cross term // uses the full dot product f.m, not the per-channel f_c m_c. return -(fixedOutput / v.unsqueeze(-1) - diff --git a/include/itkImpactBatchBudget.h b/include/itkImpactBatchBudget.h index ae9ca78..5413158 100644 --- a/include/itkImpactBatchBudget.h +++ b/include/itkImpactBatchBudget.h @@ -109,10 +109,16 @@ FitBatchToDevice(int64_t estimate, double usable, TPeakOf && peakOf) return batch; } -/** Patches of `configuration` that fit on `device`, forward and backward; 0 if not measurable - * (CPU, CPU-only LibTorch, or no strictly positive patch size). Moves the model to the device. */ +/** Patches of `configuration` that fit on `device`, forward and backward, up to `largest` if it is + * positive; 0 if not measurable (CPU, CPU-only LibTorch, or no strictly positive patch size). The + * trial batch never exceeds `largest`: on a shared GPU, a probe sized to the whole free memory + * took it from the other processes for a batch that would never be run. Moves the model to the + * device. */ inline int64_t -MeasureBatchBudget(const ImpactModelConfiguration & configuration, const torch::Device & device, double safety = 0.8) +MeasureBatchBudget(const ImpactModelConfiguration & configuration, + const torch::Device & device, + double safety = 0.8, + int64_t largest = 0) { #if ITK_IMPACT_HAS_CUDA_MEMORY_STATS const std::vector & patchSize = configuration.GetPatchSize(); @@ -160,27 +166,33 @@ MeasureBatchBudget(const ImpactModelConfiguration & configuration, const torch:: const auto stats = c10::cuda::CUDACachingAllocator::getDeviceStats(index); const int64_t cached = stats.reserved_bytes[aggregate].current - stats.allocated_bytes[aggregate].current; const double usable = static_cast(static_cast(freeBytes) + std::max(cached, 0)) * safety; - return FitBatchToDevice( - std::max(static_cast(usable / static_cast(perPatch)), 1), usable, [&](int64_t batch) { - try - { - return peakOf(batch); - } - catch (...) - { - c10::cuda::CUDACachingAllocator::emptyCache(); - throw; - } - }); + int64_t estimate = std::max(static_cast(usable / static_cast(perPatch)), 1); + if (largest > 0) + { + estimate = std::min(estimate, largest); + } + return FitBatchToDevice(estimate, usable, [&](int64_t batch) { + try + { + return peakOf(batch); + } + catch (...) + { + c10::cuda::CUDACachingAllocator::emptyCache(); + throw; + } + }); #else (void)configuration; (void)device; (void)safety; + (void)largest; return 0; #endif } -/** Bound every configuration by min(requested, measured budget); 0 for either means "no bound". */ +/** Bound every configuration by min(requested, measured budget); 0 for either means "no bound". + * The measurement tries no batch larger than `requested`. */ inline void ConfigureBatchSize(std::vector & configurations, const torch::Device & device, @@ -188,7 +200,7 @@ ConfigureBatchSize(std::vector & configurations, { for (ImpactModelConfiguration & configuration : configurations) { - const int64_t budget = MeasureBatchBudget(configuration, device); + const int64_t budget = MeasureBatchBudget(configuration, device, 0.8, requested); const int64_t bound = requested > 0 && budget > 0 ? std::min(requested, budget) : std::max(requested, budget); SetBatchSize(configuration, bound); } diff --git a/include/itkImpactCoarseRegistration.hxx b/include/itkImpactCoarseRegistration.hxx index 04d784c..74cb09b 100644 --- a/include/itkImpactCoarseRegistration.hxx +++ b/include/itkImpactCoarseRegistration.hxx @@ -28,6 +28,7 @@ #include #include +#include #include #include @@ -325,7 +326,10 @@ ImpactCoarseRegistration::GenerateData() } const torch::Device device(m_Device); - torch::manual_seed(m_Seed); + // The stage's own generator, on the CPU. torch::manual_seed reseeded every generator of the + // process, a Python caller's included, and a draw on the device also followed whatever else draws + // there, such as the batch-budget probe. Draws are made here and moved to the device. + at::Generator generator = at::make_generator(m_Seed); torch::NoGradGuard noGrad; // the coarse stage is non-differentiable // ---- Fixed and moving onto the same (fixed) voxel grid; resample only if geometries differ. ---- @@ -420,9 +424,7 @@ ImpactCoarseRegistration::GenerateData() if (kept > 0 && kept < fixedLayers[l].size(1)) { const torch::Tensor channels = - torch::randperm(fixedLayers[l].size(1), torch::TensorOptions().dtype(torch::kLong)) - .narrow(0, 0, kept) - .to(device); + torch::randperm(fixedLayers[l].size(1), generator, torch::kLong).narrow(0, 0, kept).to(device); fixedLayers[l] = fixedLayers[l].index_select(1, channels).contiguous(); movingLayers[l] = movingLayers[l].index_select(1, channels).contiguous(); } diff --git a/include/itkImpactFineRegistration.hxx b/include/itkImpactFineRegistration.hxx index 175ed85..f26f48f 100644 --- a/include/itkImpactFineRegistration.hxx +++ b/include/itkImpactFineRegistration.hxx @@ -31,6 +31,7 @@ #include #include +#include #include #include @@ -221,7 +222,10 @@ ImpactFineRegistration::GenerateData() const bool featureMode = !m_FixedModelsConfiguration.empty(); const torch::Device device(m_Device); - torch::manual_seed(m_Seed); + // The stage's own generator, on the CPU. torch::manual_seed reseeded every generator of the + // process, a Python caller's included, and a draw on the device also followed whatever else draws + // there, such as the batch-budget probe. Draws are made here and moved to the device. + at::Generator generator = at::make_generator(m_Seed); // Release the Python GIL for the whole torch section: the Adam loop runs loss.backward(), // and libtorch's Python autograd engine (active whenever we run under Python) requires the @@ -526,8 +530,7 @@ ImpactFineRegistration::GenerateData() const int64_t kept = l < m_SubsetFeatures.size() ? static_cast(m_SubsetFeatures[l]) : 0; if (kept > 0 && kept < layers[l].size(1)) { - subsets[l] = - torch::randperm(layers[l].size(1), torch::TensorOptions().dtype(torch::kLong)).narrow(0, 0, kept).to(device); + subsets[l] = torch::randperm(layers[l].size(1), generator, torch::kLong).narrow(0, 0, kept).to(device); } } }; @@ -588,11 +591,11 @@ ImpactFineRegistration::GenerateData() } count = std::max(1, std::llround(m_SamplingPercentage * static_cast(candidates))); drawn = fixedLayerCandidates[l].index_select( - 0, torch::randint(candidates, { count }, torch::TensorOptions().dtype(torch::kLong).device(device))); + 0, torch::randint(candidates, { count }, generator, torch::kLong).to(device)); } else { - drawn = torch::randint(voxels, { count }, torch::TensorOptions().dtype(torch::kLong).device(device)); + drawn = torch::randint(voxels, { count }, generator, torch::kLong).to(device); } const torch::Tensor points = layerBaseGrid[l].reshape({ voxels, ImageDimension }).index_select(0, drawn); std::vector pointShape(ImageDimension + 2, 1); // grid_sample's {1, [1,] 1, n, N} @@ -1036,7 +1039,7 @@ ImpactFineRegistration::GenerateData() voxels *= extent; } const int64_t count = std::max(1, std::llround(m_SamplingPercentage * static_cast(voxels))); - const torch::Tensor flat = torch::randint(voxels, { count }, torch::TensorOptions().dtype(torch::kLong)); + const torch::Tensor flat = torch::randint(voxels, { count }, generator, torch::kLong); torch::Tensor points = torch::empty({ count, static_cast(ImageDimension) }, torch::kLong); int64_t stride = 1; for (unsigned int d = 0; d < ImageDimension; ++d) // ITK index, x fastest in the flat voxel number diff --git a/include/itkImpactImageToImageMetricv4.h b/include/itkImpactImageToImageMetricv4.h index 022bf6d..2bab0bb 100644 --- a/include/itkImpactImageToImageMetricv4.h +++ b/include/itkImpactImageToImageMetricv4.h @@ -205,15 +205,16 @@ class ITK_TEMPLATE_EXPORT ImpactImageToImageMetricv4 * points, so lower it for large patches and raise it for small ones. Ignored in Static * mode, which interpolates a precomputed map instead of running the model. * - * The default is not a compromise between speed and memory so much as a floor. A batch of a - * few points makes tensors that are large enough for LibTorch to hand its ops to the whole - * intra-op thread pool and far too small to pay for the barriers, which measured two orders - * of magnitude slower than either a batch of one or a batch past a few dozen. Values well - * below the default are best avoided whatever the memory. + * The default, zero, chooses: on a CUDA device as many points as it can hold, which + * Initialize() measures (itkImpactBatchBudget.h) -- a GPU needs thousands to be busy; on the + * CPU 64. On the CPU a batch of a few points makes tensors that are large enough for LibTorch + * to hand its ops to the whole intra-op thread pool and far too small to pay for the barriers, + * which measured two orders of magnitude slower than either a batch of one or a batch past a + * few dozen, while a batch of hundreds falls out of the caches. Values well below 64 are best + * avoided whatever the memory. * - * On a CUDA device the batch never exceeds what the device can hold, which Initialize() - * measures (itkImpactBatchBudget.h); a batch that still fails to allocate is halved and - * replayed. Zero means that measured bound alone, or the whole work unit on the CPU. */ + * On a CUDA device a batch never exceeds what the device can hold; one that still fails to + * allocate is halved and replayed. */ itkSetMacro(BatchSize, unsigned int); itkGetConstMacro(BatchSize, unsigned int); @@ -310,11 +311,11 @@ class ITK_TEMPLATE_EXPORT ImpactImageToImageMetricv4 /** Latched at the first evaluation after Initialize(); mutable because the evaluation is const. */ mutable Impact::LossNormalization m_LossNormalization; std::string m_Device = "cpu"; - // Zero means "seed from the clock", which is what the per-work-unit generator does with it. - // It must have a value even when the user never calls SetSeed(): it is read on every - // evaluation, and an indeterminate one would make the metric unreproducible at random. + // Zero means "seed from the clock", drawn once per metric. It must have a value even when the + // user never calls SetSeed(): it is read on every evaluation, and an indeterminate one would + // make the metric unreproducible at random. unsigned int m_Seed{ 0 }; - unsigned int m_BatchSize{ 64 }; + unsigned int m_BatchSize{ 0 }; std::vector> m_features_indexes; }; diff --git a/include/itkImpactImageToImageMetricv4.hxx b/include/itkImpactImageToImageMetricv4.hxx index 650b34d..499e2a5 100644 --- a/include/itkImpactImageToImageMetricv4.hxx +++ b/include/itkImpactImageToImageMetricv4.hxx @@ -95,6 +95,10 @@ ImpactImageToImageMetricv4; this->m_DenseGetValueAndDerivativeThreader = ImpactDenseGetValueAndDerivativeThreaderType::New(); this->m_SparseGetValueAndDerivativeThreader = ImpactSparseGetValueAndDerivativeThreaderType::New(); + + // IMPACT compares features and never reads the images' gradients: with the filter on, Initialize() + // smoothed the whole moving image into a gradient image that nothing used. + this->SetUseMovingImageGradientFilter(false); } template (this->m_BatchSize > 0 ? this->m_BatchSize : this->GetNumberOfDomainPoints()); Impact::PythonGilReleaseGuard noGil; - Impact::ConfigureBatchSize(this->m_MovingModelsConfiguration, device, static_cast(this->m_BatchSize)); - Impact::ConfigureBatchSize(this->m_FixedModelsConfiguration, device, static_cast(this->m_BatchSize)); + Impact::ConfigureBatchSize(this->m_MovingModelsConfiguration, device, largest); + Impact::ConfigureBatchSize(this->m_FixedModelsConfiguration, device, largest); } else { diff --git a/include/itkImpactImageToImageMetricv4GetValueAndDerivativeThreader.h b/include/itkImpactImageToImageMetricv4GetValueAndDerivativeThreader.h index 9d8e284..8501ed7 100644 --- a/include/itkImpactImageToImageMetricv4GetValueAndDerivativeThreader.h +++ b/include/itkImpactImageToImageMetricv4GetValueAndDerivativeThreader.h @@ -23,6 +23,7 @@ #include "ImpactLoss.h" #include "itkImpactLossNormalization.h" #include +#include #include namespace itk @@ -83,22 +84,19 @@ class ITK_TEMPLATE_EXPORT ImpactImageToImageMetricv4GetValueAndDerivativeThreade AfterThreadedExecution() override; - /** Compute the local, per-point contribution to the metric value and derivative. */ + /** Never called: ThreadedExecution() evaluates the points in batches, in both modes. */ bool - ProcessPoint(const VirtualIndexType & virtualIndex, - const VirtualPointType & virtualPoint, - const FixedImagePointType & mappedFixedPoint, - const FixedImagePixelType & fixedImageValue, - const FixedImageGradientType & mappedFixedImageGradient, - const MovingImagePointType & mappedMovingPoint, - const MovingImagePixelType & movingImageValue, - const MovingImageGradientType & movingImageGradient, - MeasureType & metricValueReturn, - DerivativeType & localDerivativeReturn, - const ThreadIdType threadId) const override; - - std::vector - GetSubsetOfFeatures(const std::vector & features_index, std::mt19937 & randomGenerator, int n) const; + ProcessPoint(const VirtualIndexType &, + const VirtualPointType &, + const FixedImagePointType &, + const FixedImagePixelType &, + const FixedImageGradientType &, + const MovingImagePointType &, + const MovingImagePixelType &, + const MovingImageGradientType &, + MeasureType &, + DerivativeType &, + const ThreadIdType) const override; /** * \brief Per-work-unit accumulator of loss values and gradients, one loss object per @@ -113,22 +111,10 @@ class ITK_TEMPLATE_EXPORT ImpactImageToImageMetricv4GetValueAndDerivativeThreade std::vector m_layersWeight; SizeValueType m_numberOfPixelsCounted; int m_nb_parameters; - std::mt19937 m_randomGenerator; void - init(std::vector distance_name, - std::vector layersWeight, - unsigned int seed, - unsigned int nb_parameters) + init(std::vector distance_name, std::vector layersWeight, unsigned int nb_parameters) { - if (seed > 0) - { - this->m_randomGenerator = std::mt19937(seed); - } - else - { - this->m_randomGenerator = std::mt19937(time(nullptr)); - } this->m_layersWeight = layersWeight; for (std::string name : distance_name) { @@ -217,7 +203,7 @@ class ITK_TEMPLATE_EXPORT ImpactImageToImageMetricv4GetValueAndDerivativeThreade itkAlignedTypedef(ITK_CACHE_LINE_ALIGNMENT, PaddedLossPerThreadStruct, AlignedLossPerThreadStruct); /** A point of the virtual domain that maps inside both images, resolved into each of them. - * Online mode evaluates points in batches, so it gathers them before running any model. */ + * Both modes evaluate points in batches, so they gather them first. */ struct Sample { VirtualIndexType m_VirtualIndex; @@ -231,9 +217,9 @@ class ITK_TEMPLATE_EXPORT ImpactImageToImageMetricv4GetValueAndDerivativeThreade using PatchPlaneType = Matrix; - /** Walk this work unit's share of the domain. Online mode walks it here so it can run the - * models on a batch of points at a time; Static mode is left to the framework's own - * per-point walk, which ends in ProcessPoint(). */ + /** Walk this work unit's share of the domain and evaluate its points in batches: Static, the + * losses take a batch of feature vectors at a time instead of a few small tensor operations + * per point. Online, the work units only collect their points, for EvaluateOnline(). */ void ThreadedExecution(const DomainType & domain, const ThreadIdType threadId) override; @@ -266,13 +252,61 @@ class ITK_TEMPLATE_EXPORT ImpactImageToImageMetricv4GetValueAndDerivativeThreade torch::Tensor & movingGradients, bool computeDerivative) const; - /** Run the models on one batch of points, accumulating the layer losses and, when a - * derivative is requested, their gradient. */ + /** A batch of points as the models take it, sampled on the CPU: per model, the fixed and + * moving patches ([B, nVox]) and the moving interpolator's gradient at every patch voxel + * ([B, nVox, MovingDim]); the moving transform's Jacobian ([B, MovingDim, P]) and the global + * parameter index of each of its columns ([B, P]). */ + struct PatchBatch + { + std::vector m_FixedPatches; + std::vector m_MovingPatches; + std::vector m_MovingGradients; + torch::Tensor m_TransformJacobian; + torch::Tensor m_NonZeroJacobianIndices; + }; + + /** Cut the patches of a batch of points, and with \p derivative the moving gradients and the + * transform's Jacobian at each; over all the threads when \p parallel, which a thread already + * running its own share leaves off. */ + PatchBatch + SampleBatch(typename std::vector::const_iterator first, + typename std::vector::const_iterator last, + bool parallel, + bool derivative) const; + + /** A batch of patch values, [B, nVox], as model \p model of the moving or the fixed image takes + * them: [B, C, s...] on the device. */ + torch::Tensor + ModelPatches(const torch::Tensor & values, size_t model, bool moving) const; + + /** Run model \p model of the moving or the fixed image on \p patches: the central feature vector + * of every kept layer, at the channels its comparison samples, [B, C] each on the device. + * \p comparison is the first of the model's comparisons. */ + std::vector + Features(size_t model, bool moving, const torch::Tensor & patches, size_t comparison) const; + + /** d(loss)/d(parameter) of a batch at every comparison, [B, 1, P] each, to be added at the + * batch's NonZeroJacobianIndices. The moving models run with gradient, and each comparison's + * backward pass is seeded with `modulator(comparison, movingFeatures)`: d(loss)/d(feature) at + * the batch's points, [B, C]. */ + std::vector + Backpropagate(const PatchBatch & batch, + const std::function & modulator) const; + + /** Online mode: run the models on the points every work unit collected. */ + void + EvaluateOnline(); + + /** Read the feature maps at one batch of points, accumulating the layer losses and, when a + * derivative is requested, their gradient (Static mode). */ void - EvaluateJacobianBatch(typename std::vector::const_iterator first, - typename std::vector::const_iterator last, - const ThreadIdType threadId, - LossPerThreadStruct & loss) const; + EvaluateStaticBatch(typename std::vector::const_iterator first, + typename std::vector::const_iterator last, + const ThreadIdType threadId, + LossPerThreadStruct & loss) const; + + /** The points each work unit collected, online. */ + std::vector> m_OnlineSamples; /** One accumulator per work unit, allocated in BeforeThreadedExecution. */ mutable std::unique_ptr m_LossThreadStruct{ nullptr }; @@ -287,6 +321,13 @@ class ITK_TEMPLATE_EXPORT ImpactImageToImageMetricv4GetValueAndDerivativeThreade * than once per evaluation. Shared by every work unit, so the value cannot depend on how the * domain was partitioned. */ unsigned int m_PlaneSeed{ 0 }; + + /** The channels every point of an evaluation compares, per comparison, and the same as index + * tensors on the device for the online mode. Drawn once in BeforeThreadedExecution from + * m_PlaneSeed: a loss such as NCC accumulates its statistics channel by channel over the points, + * so all of them have to compare the same channels, whatever the partition or the batches. */ + std::vector> m_SubsetsOfFeatures; + std::vector m_SubsetsOfFeaturesOnDevice; }; } // end namespace itk diff --git a/include/itkImpactImageToImageMetricv4GetValueAndDerivativeThreader.hxx b/include/itkImpactImageToImageMetricv4GetValueAndDerivativeThreader.hxx index 85de8dd..4800fa5 100644 --- a/include/itkImpactImageToImageMetricv4GetValueAndDerivativeThreader.hxx +++ b/include/itkImpactImageToImageMetricv4GetValueAndDerivativeThreader.hxx @@ -74,17 +74,40 @@ ImpactImageToImageMetricv4GetValueAndDerivativeThreader(numWorkUnitsUsed); + this->m_OnlineSamples.assign(numWorkUnitsUsed, {}); for (ThreadIdType i = 0; i < numWorkUnitsUsed; ++i) { - // The feature subset follows the same seed as the planes above: with Seed set it is that - // seed, with Seed zero it is the clock drawn once. Reseeding from the clock here on every - // evaluation made the subset change across a second boundary, so two evaluations of the - // same parameters did not return the same value. - this->m_LossThreadStruct[i].init(this->m_ImpactAssociate->GetDistance(), - this->m_ImpactAssociate->GetLayersWeight(), - this->m_PlaneSeed, - globalDerivativeSize); + this->m_LossThreadStruct[i].init( + this->m_ImpactAssociate->GetDistance(), this->m_ImpactAssociate->GetLayersWeight(), globalDerivativeSize); + } + + // The feature subsets follow the same seed as the planes above, drawn afresh from it on every + // evaluation, so two evaluations of the same parameters compare the same channels. + std::mt19937 generator(this->m_PlaneSeed); + const torch::Device device(this->m_ImpactAssociate->GetDevice()); + this->m_SubsetsOfFeatures = this->m_ImpactAssociate->m_features_indexes; + this->m_SubsetsOfFeaturesOnDevice.clear(); + for (size_t comparison = 0; comparison < this->m_SubsetsOfFeatures.size(); ++comparison) + { + std::vector & subset = this->m_SubsetsOfFeatures[comparison]; + const unsigned int kept = this->m_ImpactAssociate->GetSubsetFeatures()[comparison]; + if (kept < subset.size()) + { + std::shuffle(subset.begin(), subset.end(), generator); + subset.resize(kept); + } + this->m_SubsetsOfFeaturesOnDevice.push_back( + torch::tensor(std::vector(subset.begin(), subset.end()), torch::kInt64).to(device)); + } + // Static mode reads exactly these channels at every point: have the maps lay them side by side. + auto & internals = *this->m_ImpactAssociate->m_Internals; + for (size_t comparison = 0; comparison < internals.m_FixedFeaturesMaps.size(); ++comparison) + { + internals.m_FixedFeaturesMaps[comparison].m_FeaturesMapInterpolator->SetActiveChannels( + this->m_SubsetsOfFeatures[comparison]); + internals.m_MovingFeaturesMaps[comparison].m_FeaturesMapInterpolator->SetActiveChannels( + this->m_SubsetsOfFeatures[comparison]); } } @@ -98,6 +121,11 @@ ImpactImageToImageMetricv4GetValueAndDerivativeThreaderm_ImpactAssociate->GetMode() == "Jacobian") + { + this->EvaluateOnline(); + } + // Reduce the per-work-unit loss accumulators into work unit 0. const ThreadIdType numWorkUnitsUsed = this->GetNumberOfWorkUnitsUsed(); for (ThreadIdType i = 1; i < numWorkUnitsUsed; ++i) @@ -127,85 +155,47 @@ ImpactImageToImageMetricv4GetValueAndDerivativeThreader -std::vector -ImpactImageToImageMetricv4GetValueAndDerivativeThreader:: - GetSubsetOfFeatures(const std::vector & features_index, std::mt19937 & randomGenerator, int n) const -{ - if (features_index.size() == static_cast(n)) - return features_index; - - std::vector shuffled = features_index; - std::shuffle(shuffled.begin(), shuffled.end(), randomGenerator); - shuffled.resize(n); - return shuffled; -} - template void ImpactImageToImageMetricv4GetValueAndDerivativeThreader:: ThreadedExecution(const DomainType & domain, const ThreadIdType threadId) { - // Static mode reads a precomputed feature map, so a point costs one interpolation and the - // framework's own walk -- which ends in ProcessPoint() -- is the right granularity. - if (this->m_ImpactAssociate->GetMode() != "Jacobian") - { - Superclass::ThreadedExecution(domain, threadId); - return; - } + // Gather this work unit's points first and evaluate them in batches, which is the shape the + // elastix metric has: there the sampler hands its whole point set over in a single call. + // Online, walking point by point would cost one forward and one backward pass each; Static, a + // few small tensor operations each, and the framework's walk also computed an image gradient + // per point that IMPACT never reads. + std::vector samples = this->CollectSamples(domain); + LossPerThreadStruct & loss = this->m_LossThreadStruct[threadId]; - // Online mode instead runs the network on a patch cut around every point, so walking point - // by point costs one forward and one backward pass each. Gather this work unit's points - // first and push them through the models in batches, which is the shape the elastix metric - // has: there the sampler hands its whole point set over in a single call. - const std::vector samples = this->CollectSamples(domain); - LossPerThreadStruct & loss = this->m_LossThreadStruct[threadId]; + // ProcessVirtualPoint() would have kept these counts and we bypassed it: the first is what + // the base metric reduces into GetNumberOfValidPoints(), the second what normalizes the loss. + this->m_GetValueAndDerivativePerThreadVariables[threadId].NumberOfValidPoints += samples.size(); + loss.m_numberOfPixelsCounted += samples.size(); - // BatchSize capped by the device budget (itkImpactBatchBudget.h); 0 = the budget alone, or - // the whole work unit on the CPU. A batch that runs out of memory is halved and replayed - // after the losses have been rewound. - const auto & movingConfigs = this->m_ImpactAssociate->GetMovingModelsConfiguration(); - const torch::Device device(this->m_ImpactAssociate->GetDevice()); - const auto requested = static_cast(this->m_ImpactAssociate->GetBatchSize()); - for (auto first = samples.begin(); first != samples.end();) + if (this->m_ImpactAssociate->GetMode() == "Jacobian") { - const int64_t bound = Impact::EffectiveBatchSize(movingConfigs); - int64_t batch = requested > 0 ? requested : static_cast(samples.size()); - batch = std::max(bound > 0 ? std::min(batch, bound) : batch, 1); - const auto last = first + std::min(batch, samples.end() - first); - - std::vector before; - for (const auto & l : loss.m_losses) + this->m_OnlineSamples[threadId] = std::move(samples); // for EvaluateOnline() + } + else + { + // A batch holds a [batch, channels, parameters] Jacobian: about 4 MB per work unit, so that a + // dense domain under a B-spline transform stays bounded. + size_t channels = 1; + for (const std::vector & subset : this->m_SubsetsOfFeatures) { - before.push_back(l->SaveState()); + channels = std::max(channels, subset.size()); } - try + const size_t parameters = this->GetComputeDerivative() ? this->GetCachedNumberOfLocalParameters() : 1; + const auto batch = static_cast(std::max(1, (size_t{ 1 } << 20) / (channels * parameters))); + for (auto first = samples.begin(); first != samples.end();) { - this->EvaluateJacobianBatch(first, last, threadId, loss); + const auto last = first + std::min(batch, samples.end() - first); + this->EvaluateStaticBatch(first, last, threadId, loss); first = last; } - catch (const std::exception & error) - { - if (!Impact::IsDeviceOutOfMemory(error) || last - first <= 1) - { - throw; - } - for (size_t l = 0; l < loss.m_losses.size(); ++l) - { - loss.m_losses[l]->RestoreState(before[l]); - } - for (const ImpactModelConfiguration & config : movingConfigs) - { - Impact::ShrinkBatchSize(config, last - first, device); - } - } } - // ProcessVirtualPoint() would have kept these counts and we bypassed it: the first is what - // the base metric reduces into GetNumberOfValidPoints(), the second what normalizes the loss. - this->m_GetValueAndDerivativePerThreadVariables[threadId].NumberOfValidPoints += samples.size(); - loss.m_numberOfPixelsCounted += samples.size(); - this->m_Associate->FinalizeThread(threadId); } @@ -424,348 +414,498 @@ ImpactImageToImageMetricv4GetValueAndDerivativeThreader -void +auto ImpactImageToImageMetricv4GetValueAndDerivativeThreader:: - EvaluateJacobianBatch(typename std::vector::const_iterator first, - typename std::vector::const_iterator last, - const ThreadIdType threadId, - LossPerThreadStruct & loss) const + SampleBatch(typename std::vector::const_iterator first, + typename std::vector::const_iterator last, + bool parallel, + bool derivative) const -> PatchBatch { + using JacobianType = typename TImageToImageMetric::JacobianType; + using MovingTransformType = typename TImageToImageMetric::MovingTransformType; constexpr unsigned int movingDimension = ImageToImageMetricv4Type::MovingImageDimension; - const auto batchSize = static_cast(std::distance(first, last)); - const bool computeDerivative = this->GetComputeDerivative(); - const torch::Device device(this->m_ImpactAssociate->GetDevice()); - std::unique_lock deviceLock(Impact::DeviceForwardMutex(), std::defer_lock); + const auto size = static_cast(last - first); + const auto & configs = this->m_ImpactAssociate->GetFixedModelsConfiguration(); - // d(moving coordinate)/d(parameter) at every point of the batch, and the global parameter - // index each of its columns belongs to. - torch::Tensor transformJacobian; // [B, MovingDim, P] - torch::Tensor nonZeroJacobianIndices; // [B, P] - if (computeDerivative) + PatchBatch batch; + for (const ImpactModelConfiguration & config : configs) { - using JacobianType = typename TImageToImageMetric::JacobianType; - using MovingTransformType = typename TImageToImageMetric::MovingTransformType; - JacobianType & jacobian = this->m_GetValueAndDerivativePerThreadVariables[threadId].MovingTransformJacobian; - JacobianType & jacobianPositional = - this->m_GetValueAndDerivativePerThreadVariables[threadId].MovingTransformJacobianPositional; + const std::vector & patchSize = config.GetPatchSize(); + const int64_t voxels = std::accumulate(patchSize.begin(), patchSize.end(), int64_t{ 1 }, std::multiplies<>()); + batch.m_FixedPatches.push_back(torch::zeros({ size, voxels }, torch::kFloat32)); + batch.m_MovingPatches.push_back(torch::zeros({ size, voxels }, torch::kFloat32)); + batch.m_MovingGradients.push_back(torch::zeros({ size, voxels, movingDimension }, torch::kFloat32)); + } - const NumberOfParametersType parameterCount = this->GetCachedNumberOfLocalParameters(); - transformJacobian = torch::empty( - { batchSize, static_cast(movingDimension), static_cast(parameterCount) }, torch::kFloat32); - nonZeroJacobianIndices = torch::empty({ batchSize, static_cast(parameterCount) }, torch::kInt64); - auto jacobianAccessor = transformJacobian.accessor(); - auto indexAccessor = nonZeroJacobianIndices.accessor(); + // d(moving coordinate)/d(parameter) at every point, and the global parameter index of each of + // its columns. A global transform (affine, rigid) has offset 0; a local-support displacement + // field only affects the parameter block of the point's own virtual voxel. + const auto parameters = static_cast(this->GetCachedNumberOfLocalParameters()); + const auto transform = this->m_ImpactAssociate->GetMovingTransform(); + const bool field = transform->GetTransformCategory() == MovingTransformType::TransformCategoryEnum::DisplacementField; + float * jacobianOut = nullptr; + int64_t * indicesOut = nullptr; + if (derivative) + { + batch.m_TransformJacobian = torch::empty({ size, movingDimension, parameters }, torch::kFloat32); + batch.m_NonZeroJacobianIndices = torch::empty({ size, parameters }, torch::kInt64); + jacobianOut = batch.m_TransformJacobian.template data_ptr(); + indicesOut = batch.m_NonZeroJacobianIndices.template data_ptr(); + } - for (int64_t row = 0; row < batchSize; ++row) + auto sampleRows = [&](int64_t begin, int64_t end) { + JacobianType jacobian(movingDimension, parameters); + JacobianType jacobianPositional(movingDimension, movingDimension); + for (int64_t row = begin; row < end; ++row) { const Sample & sample = *(first + row); - this->m_ImpactAssociate->GetMovingTransform()->ComputeJacobianWithRespectToParametersCachedTemporaries( + for (size_t i = 0; i < configs.size(); ++i) + { + this->SamplePatch( + sample, i, row, batch.m_FixedPatches[i], batch.m_MovingPatches[i], batch.m_MovingGradients[i], derivative); + } + if (!derivative) + { + continue; + } + transform->ComputeJacobianWithRespectToParametersCachedTemporaries( sample.m_VirtualPoint, jacobian, jacobianPositional); - for (unsigned int dimension = 0; dimension < movingDimension; ++dimension) + for (unsigned int d = 0; d < movingDimension; ++d) { - for (NumberOfParametersType parameter = 0; parameter < parameterCount; ++parameter) + for (int64_t parameter = 0; parameter < parameters; ++parameter) { - jacobianAccessor[row][dimension][parameter] = static_cast(jacobian(dimension, parameter)); + jacobianOut[(row * movingDimension + d) * parameters + parameter] = + static_cast(jacobian(d, parameter)); } } - // Map local parameter indices to global ones. A global transform (affine, rigid) has - // offset 0; a local-support displacement field only affects the parameter block of this - // point's own virtual voxel. - int64_t parameterOffset = 0; - if (this->m_ImpactAssociate->GetMovingTransform()->GetTransformCategory() == - MovingTransformType::TransformCategoryEnum::DisplacementField) + const int64_t offset = + field ? static_cast( + this->m_ImpactAssociate->ComputeParameterOffsetFromVirtualIndex(sample.m_VirtualIndex, parameters)) + : 0; + for (int64_t parameter = 0; parameter < parameters; ++parameter) { - parameterOffset = static_cast( - this->m_ImpactAssociate->ComputeParameterOffsetFromVirtualIndex(sample.m_VirtualIndex, parameterCount)); - } - for (NumberOfParametersType parameter = 0; parameter < parameterCount; ++parameter) - { - indexAccessor[row][parameter] = parameterOffset + static_cast(parameter); + indicesOut[row * parameters + parameter] = offset + parameter; } } + }; + if (!parallel) + { + sampleRows(0, size); + return batch; } + // The rows are independent: spread them over the threads, idle between two evaluations. + MultiThreaderBase * const threader = this->GetMultiThreader(); + const auto parts = std::min(size, threader->GetNumberOfWorkUnits()); + threader->ParallelizeArray( + 0, parts, [&](SizeValueType part) { sampleRows(size * part / parts, size * (part + 1) / parts); }, nullptr); + return batch; +} + + +template +torch::Tensor +ImpactImageToImageMetricv4GetValueAndDerivativeThreader:: + ModelPatches(const torch::Tensor & values, size_t model, bool moving) const +{ + // [B, C, s...]. PatchTensorShape is the one place that knows the spatial extents are reversed + // (image axis 0 innermost): SamplePatch fills each row with image axis 0 varying fastest, which + // is the layout ImageToTensorFilter produces and the one the TorchScript models are trained on. + // The single intensity channel is then repeated to the channel count the model expects. Both + // images are cut to the fixed configuration's patch. + const ImpactModelConfiguration & fixedConfig = this->m_ImpactAssociate->GetFixedModelsConfiguration()[model]; + const ImpactModelConfiguration & config = + moving ? this->m_ImpactAssociate->GetMovingModelsConfiguration()[model] : fixedConfig; + const std::vector spatial = Impact::PatchTensorShape(fixedConfig); + std::vector shape{ values.size(0), 1 }; + shape.insert(shape.end(), spatial.begin(), spatial.end()); + std::vector channels(shape.size(), 1); + channels[1] = static_cast(fixedConfig.GetNumberOfChannels()); + return values.reshape(shape) + .repeat(channels) + .to(torch::Device(this->m_ImpactAssociate->GetDevice())) + .to(GetModelDtype(config)); +} + + +template +std::vector +ImpactImageToImageMetricv4GetValueAndDerivativeThreader:: + Features(size_t model, bool moving, const torch::Tensor & patches, size_t comparison) const +{ + // One configuration serves both images, so each image's metadata travels with its patches. + const auto & internals = *this->m_ImpactAssociate->m_Internals; + const auto & config = moving ? this->m_ImpactAssociate->GetMovingModelsConfiguration()[model] + : this->m_ImpactAssociate->GetFixedModelsConfiguration()[model]; + std::vector outputs = + Forward(config, patches, moving ? internals.m_MovingImageMetadata : internals.m_FixedImageMetadata); // The layers the metric compares, in mask order: one feature map, one loss and one weight // each, appended to the same flat list across models that Static builds. - auto keptLayers = [](std::vector & outputs, const ImpactModelConfiguration & config) { - std::vector layers; - const auto & mask = config.GetLayersMask(); - for (size_t it = 0; it < mask.size(); ++it) + const std::vector & mask = config.GetLayersMask(); + std::vector features; + for (size_t it = 0; it < mask.size(); ++it) + { + if (!mask[it]) { - if (mask[it]) - { - layers.push_back(outputs[it].toTensor()); - } + continue; } - return layers; - }; - // The central feature vector of every point of a layer tensor [B, C, s...] -> [B, C]. - auto centerFeatures = [](torch::Tensor layer) { + // The central feature vector of every point: [B, C, s...] -> [B, C]. + torch::Tensor layer = outputs[it].toTensor(); while (layer.dim() > 2) { layer = layer.select(2, layer.size(2) / 2); } - return layer; - }; + features.push_back(layer.index_select(1, this->m_SubsetsOfFeaturesOnDevice[comparison++]).to(torch::kFloat32)); + } + return features; +} - const auto & fixedConfigs = this->m_ImpactAssociate->GetFixedModelsConfiguration(); - const auto & movingConfigs = this->m_ImpactAssociate->GetMovingModelsConfiguration(); - // One configuration serves both images, so each image's metadata travels with its patches. - const ImpactImageMetadata & fixedMetadata = this->m_ImpactAssociate->m_Internals->m_FixedImageMetadata; - const ImpactImageMetadata & movingMetadata = this->m_ImpactAssociate->m_Internals->m_MovingImageMetadata; - size_t comparison = 0; // flat index over (model, kept layer), the one the losses are read at - for (size_t i = 0; i < fixedConfigs.size(); ++i) +template +std::vector +ImpactImageToImageMetricv4GetValueAndDerivativeThreader:: + Backpropagate(const PatchBatch & batch, const std::function & modulator) const +{ + const torch::Device device(this->m_ImpactAssociate->GetDevice()); + std::vector jacobians; // one per comparison, in order + for (size_t i = 0; i < batch.m_MovingPatches.size(); ++i) { - const ImpactModelConfiguration & fixedConfig = fixedConfigs[i]; - const ImpactModelConfiguration & movingConfig = movingConfigs[i]; - const auto channels = static_cast(fixedConfig.GetNumberOfChannels()); - - const std::vector & patchSize = fixedConfig.GetPatchSize(); - const int64_t voxelCount = std::accumulate(patchSize.begin(), patchSize.end(), int64_t{ 1 }, std::multiplies<>()); - - torch::Tensor fixedPatches = torch::zeros({ batchSize, voxelCount }, torch::kFloat32); - torch::Tensor movingPatches = torch::zeros({ batchSize, voxelCount }, torch::kFloat32); - torch::Tensor movingGradients = - torch::zeros({ batchSize, voxelCount, static_cast(movingDimension) }, torch::kFloat32); - for (int64_t row = 0; row < batchSize; ++row) + // The patch values are the autograd leaf, as in the elastix metric: differentiating the + // features with respect to them and contracting with the moving interpolator's spatial + // gradients gives d(feature)/d(moving coordinate) for the whole batch at once. One forward + // pass serves every kept layer, so the graph has to survive until the last of them is done. + const torch::Tensor patches = + this->ModelPatches(batch.m_MovingPatches[i], i, true).detach().set_requires_grad(true); + const std::vector layers = this->Features(i, true, patches, jacobians.size()); + // d(patch value)/d(moving coordinate), tiled once per repeated channel so that it lines up + // with the flattened gradient below: summing over the copies sums over the channel axis, + // which is exactly the derivative with respect to the underlying intensity. + const torch::Tensor patchGradients = + batch.m_MovingGradients[i].to(device).repeat({ 1, patches.size(1), 1 }); // [B, C * nVox, MovingDim] + for (size_t k = 0; k < layers.size(); ++k) { - this->SamplePatch(*(first + row), i, row, fixedPatches, movingPatches, movingGradients, computeDerivative); + torch::Tensor features = layers[k].detach().to(torch::kCPU); + const torch::Tensor seed = modulator(jacobians.size(), features).to(device).to(layers[k].dtype()); + const torch::Tensor gradient = + torch::autograd::grad({ layers[k] }, { patches }, { seed }, k + 1 < layers.size(), false)[0]; + // d(.)/d(moving coordinate), [B, 1, MovingDim], and then d(.)/d(parameter), [B, 1, P]. + const torch::Tensor coordinate = torch::bmm(gradient.flatten(1).unsqueeze(1).to(torch::kFloat32), patchGradients); + jacobians.push_back(torch::bmm(coordinate.to(torch::kCPU), batch.m_TransformJacobian)); } + } + return jacobians; +} - // [B, C, s...]. PatchTensorShape is the one place that knows the spatial extents are - // reversed (image axis 0 innermost): SamplePatch fills each row with image axis 0 varying - // fastest, which is the layout ImageToTensorFilter produces and the one the TorchScript - // models are trained on. The single intensity channel is then repeated to the channel - // count the model expects. - const std::vector tensorShape = Impact::PatchTensorShape(fixedConfig); - std::vector patchShape{ batchSize, 1 }; - patchShape.insert(patchShape.end(), tensorShape.begin(), tensorShape.end()); - std::vector channelRepeat(patchShape.size(), 1); - channelRepeat[1] = channels; - - auto toModelPatch = [&](const torch::Tensor & values, const ImpactModelConfiguration & config) { - return values.reshape(patchShape).repeat(channelRepeat).to(device).to(GetModelDtype(config)); - }; - // The feature vector of one kept layer, restricted to the channels this comparison samples. - auto featuresOf = [&](const torch::Tensor & layer, const torch::Tensor & subsetIndices) { - return centerFeatures(layer).index_select(1, subsetIndices).to(torch::kFloat32); - }; - // Which channels of comparison `a` are compared. Drawn once per batch and per comparison, - // from the work unit's own generator. - auto subsetIndicesOf = [&](size_t a) { - const std::vector subsetOfFeatures = - this->GetSubsetOfFeatures(this->m_ImpactAssociate->m_features_indexes[a], - loss.m_randomGenerator, - this->m_ImpactAssociate->GetSubsetFeatures()[a]); - return torch::tensor(std::vector(subsetOfFeatures.begin(), subsetOfFeatures.end()), torch::kInt64) - .to(device); - }; - - // One work unit at a time on the device, so the budget holds whatever the thread count. - if (device.is_cuda() && !deviceLock.owns_lock()) - { - deviceLock.lock(); - } - // No gradient ever flows through the fixed image. - std::vector fixedLayers; +template +void +ImpactImageToImageMetricv4GetValueAndDerivativeThreader:: + EvaluateOnline() +{ + std::vector samples; + for (const std::vector & unit : this->m_OnlineSamples) + { + samples.insert(samples.end(), unit.begin(), unit.end()); + } + this->m_OnlineSamples.clear(); + const auto count = static_cast(samples.size()); + if (count == 0) + { + return; + } + + const torch::Device device(this->m_ImpactAssociate->GetDevice()); + const auto & movingConfigs = this->m_ImpactAssociate->GetMovingModelsConfiguration(); + const auto requested = static_cast(this->m_ImpactAssociate->GetBatchSize()); + const bool computeDerivative = this->GetComputeDerivative(); + auto & losses = this->m_LossThreadStruct[0].m_losses; + + // Run `evaluate(begin, end, parallel, loss)` over the samples a batch at a time, `loss` being the + // accumulator of the thread that runs it. + auto inBatches = [&](const auto & evaluate) { + if (!device.is_cuda()) { - torch::NoGradGuard noGrad; - auto outputs = Forward(fixedConfig, toModelPatch(fixedPatches, fixedConfig), fixedMetadata); - fixedLayers = keptLayers(outputs, fixedConfig); + // On the CPU the threads are the parallelism: each runs the models on its own share, + // BatchSize points at a time (0: 64, see SetBatchSize()). + const int64_t batch = requested > 0 ? requested : 64; + const auto parts = std::min(count, this->GetNumberOfWorkUnitsUsed()); + this->GetMultiThreader()->ParallelizeArray( + 0, + parts, + [&](SizeValueType part) { + const auto p = static_cast(part); + for (int64_t begin = count * p / parts, end = count * (p + 1) / parts; begin < end; begin += batch) + { + evaluate(begin, std::min(begin + batch, end), false, this->m_LossThreadStruct[part]); + } + }, + nullptr); + return; } - const auto layerCount = static_cast(fixedLayers.size()); - - if (!computeDerivative) + // On a GPU one batch at a time, as large as it holds, its patches cut by all the threads: + // BatchSize capped by the device budget (itkImpactBatchBudget.h), 0 = the budget alone. A + // batch that runs out of memory is halved and replayed after the losses have been rewound. + std::lock_guard lock(Impact::DeviceForwardMutex()); + for (int64_t begin = 0; begin < count;) { - torch::NoGradGuard noGrad; - auto outputs = Forward(movingConfig, toModelPatch(movingPatches, movingConfig), movingMetadata); - const std::vector movingLayers = keptLayers(outputs, movingConfig); - for (size_t layer = 0; layer < layerCount; ++layer, ++comparison) + const int64_t bound = Impact::EffectiveBatchSize(movingConfigs); + int64_t batch = requested > 0 ? requested : count; + batch = std::max(bound > 0 ? std::min(batch, bound) : batch, 1); + const int64_t end = std::min(begin + batch, count); + + std::vector before; + for (const auto & l : losses) { - const torch::Tensor subsetIndices = subsetIndicesOf(comparison); - torch::Tensor fixedFeatures = featuresOf(fixedLayers[layer], subsetIndices).to(torch::kCPU); - torch::Tensor movingFeatures = featuresOf(movingLayers[layer], subsetIndices).to(torch::kCPU); - loss.m_losses[comparison]->updateValue(fixedFeatures, movingFeatures); + before.push_back(l->SaveState()); + } + try + { + evaluate(begin, end, true, this->m_LossThreadStruct[0]); + begin = end; + } + catch (const std::exception & error) + { + if (!Impact::IsDeviceOutOfMemory(error) || end - begin <= 1) + { + throw; + } + for (size_t l = 0; l < losses.size(); ++l) + { + losses[l]->RestoreState(before[l]); + } + for (const ImpactModelConfiguration & config : movingConfigs) + { + Impact::ShrinkBatchSize(config, end - begin, device); + } } - continue; } - - // The patch values are the autograd leaf, as in the elastix metric: differentiating the - // features with respect to them and contracting with the moving interpolator's spatial - // gradients gives d(feature)/d(moving coordinate) for the whole batch at once. One forward - // pass serves every kept layer, so the graph has to survive until the last of them is done. - torch::Tensor movingPatch = toModelPatch(movingPatches, movingConfig).detach().set_requires_grad(true); - auto movingOutputs = Forward(movingConfig, movingPatch, movingMetadata); - const std::vector movingLayers = keptLayers(movingOutputs, movingConfig); - // d(patch value)/d(moving coordinate), tiled once per repeated channel so that it lines up - // with the flattened gradient below: summing over the copies sums over the channel axis, - // which is exactly the derivative with respect to the underlying intensity. - torch::Tensor patchGradients = movingGradients.to(device).repeat({ 1, channels, 1 }); // [B, C * nVox, MovingDim] - - // ... and then into d(.)/d(parameter), which is what the losses accumulate. - auto toParameterJacobian = [&](const torch::Tensor & featureCoordinateJacobian) { - return torch::bmm(featureCoordinateJacobian.to(torch::kCPU).to(torch::kFloat32), transformJacobian); - }; - - for (size_t layer = 0; layer < layerCount; ++layer, ++comparison) + }; + auto sample = [&](int64_t begin, int64_t end, bool parallel, bool derivative) { + return this->SampleBatch(samples.begin() + begin, samples.begin() + end, parallel, derivative); + }; + // The feature vectors of one image of a batch at every comparison, [B, C] each on the CPU. + auto featuresOf = [&](const PatchBatch & batch, bool moving) { + torch::NoGradGuard noGrad; + std::vector features; + for (size_t i = 0; i < movingConfigs.size(); ++i) { - const torch::Tensor subsetIndices = subsetIndicesOf(comparison); - torch::Tensor fixedFeatures = featuresOf(fixedLayers[layer], subsetIndices).to(torch::kCPU); - torch::Tensor movingFeatures = featuresOf(movingLayers[layer], subsetIndices); // [B, C], on the device - torch::Tensor movingFeaturesOnCpu = movingFeatures.detach().to(torch::kCPU); - const bool laterLayersFollow = layer + 1 < layerCount; - - // One backward pass over the batch, seeded by `seed`, contracted into d(.)/d(coordinate). - auto backward = [&](const torch::Tensor & seed, bool retainGraph) { - const torch::Tensor gradient = - torch::autograd::grad({ movingFeatures }, { movingPatch }, { seed }, retainGraph, false)[0]; - return torch::bmm(gradient.flatten(1).unsqueeze(1).to(torch::kFloat32), patchGradients); // [B, 1, MovingDim] - }; - - if (loss.m_losses[comparison]->IsPerPointMean()) + const torch::Tensor & values = moving ? batch.m_MovingPatches[i] : batch.m_FixedPatches[i]; + for (const torch::Tensor & layer : + this->Features(i, moving, this->ModelPatches(values, i, moving), features.size())) { - // The loss is a mean of independent per-point terms, so its gradient with respect to the - // features is known before any backward pass: seeding the pass with it yields - // d(loss)/d(moving coordinate) directly, in ONE pass for the whole batch. This is what - // the elastix metric does. - torch::Tensor modulator = - loss.m_losses[comparison]->updateValueAndGetGradientModulator(fixedFeatures, movingFeaturesOnCpu); - torch::Tensor jacobian = - toParameterJacobian(backward(modulator.to(device).to(movingFeatures.dtype()), laterLayersFollow)); - loss.m_losses[comparison]->updateDerivativeInJacobianMode(jacobian, nonZeroJacobianIndices); + features.push_back(layer.to(torch::kCPU)); } - else + } + return features; + }; + + const bool global = + computeDerivative && std::any_of(losses.begin(), losses.end(), [](const auto & l) { return !l->IsPerPointMean(); }); + if (!global) + { + inBatches([&](int64_t begin, int64_t end, bool parallel, LossPerThreadStruct & loss) { + const PatchBatch batch = sample(begin, end, parallel, computeDerivative); + std::vector fixed = featuresOf(batch, false); + if (!computeDerivative) { - // A global statistic (NCC) mixes the points together: its gradient with respect to the - // features is only known once every point has been seen, so it cannot seed the pass -- - // seeded per batch it would measure the correlation within that batch alone, which at one - // point is a zero variance and hence an identically zero gradient. Recover the full - // per-channel Jacobian instead -- one pass per kept channel, each still over the whole - // batch -- and let the loss fold it into the statistics it accumulates itself. That is - // the closed form Static mode uses, and it stays exact whatever the batch size and - // however the work units split the domain. - const int64_t keptChannels = movingFeatures.size(1); - std::vector perChannel; - perChannel.reserve(static_cast(keptChannels)); - torch::Tensor seed = torch::zeros_like(movingFeatures); - for (int64_t channel = 0; channel < keptChannels; ++channel) + std::vector moving = featuresOf(batch, true); + for (size_t c = 0; c < fixed.size(); ++c) { - seed.zero_(); - seed.select(1, channel).fill_(1.0); - perChannel.push_back(backward(seed, channel + 1 < keptChannels || laterLayersFollow).squeeze(1)); + loss.m_losses[c]->updateValue(fixed[c], moving[c]); } - torch::Tensor jacobian = toParameterJacobian(torch::stack(perChannel, 1)); // [B, C, MovingDim] -> [B, C, P] - loss.m_losses[comparison]->updateValueAndDerivativeInStaticMode( - fixedFeatures, movingFeaturesOnCpu, jacobian, nonZeroJacobianIndices); + return; + } + // A mean of per-point terms: its d(loss)/d(feature) at a point is known from that point + // alone, so the loss gives it with the batch's value, and one backward pass a comparison + // carries it to the parameters. This is what the elastix metric does. + std::vector jacobians = this->Backpropagate(batch, [&](size_t c, torch::Tensor & moving) { + return loss.m_losses[c]->updateValueAndGetGradientModulator(fixed[c], moving); + }); + torch::Tensor indices = batch.m_NonZeroJacobianIndices; + for (size_t c = 0; c < jacobians.size(); ++c) + { + loss.m_losses[c]->updateDerivativeInJacobianMode(jacobians[c], indices); } + }); + return; + } + + // A global statistic (NCC) mixes the points: d(loss)/d(feature) at one of them is only known + // once all have been seen. A first pass, without gradient, gathers the features of every point, + // from which each loss takes its value and that derivative, [N, C]; a second carries it to the + // parameters as for a per-point loss, instead of the one backward pass per channel the closed + // form Static uses would take. + std::vector fixedFeatures, movingFeatures; + for (const torch::Tensor & subset : this->m_SubsetsOfFeaturesOnDevice) + { + fixedFeatures.push_back(torch::empty({ count, subset.numel() }, torch::kFloat32)); + movingFeatures.push_back(torch::empty({ count, subset.numel() }, torch::kFloat32)); + } + inBatches([&](int64_t begin, int64_t end, bool parallel, LossPerThreadStruct &) { + const PatchBatch batch = sample(begin, end, parallel, false); + const std::vector fixed = featuresOf(batch, false); + const std::vector moving = featuresOf(batch, true); + for (size_t c = 0; c < fixed.size(); ++c) + { + fixedFeatures[c].narrow(0, begin, end - begin).copy_(fixed[c]); + movingFeatures[c].narrow(0, begin, end - begin).copy_(moving[c]); } + }); + std::vector modulators; + for (size_t c = 0; c < losses.size(); ++c) + { + modulators.push_back(losses[c]->updateValueAndGetGradientModulator(fixedFeatures[c], movingFeatures[c])); } + std::mutex accumulating; // the CPU threads all add to work unit 0's derivative + inBatches([&](int64_t begin, int64_t end, bool parallel, LossPerThreadStruct &) { + const PatchBatch batch = sample(begin, end, parallel, true); + std::vector jacobians = this->Backpropagate( + batch, [&](size_t c, torch::Tensor &) { return modulators[c].narrow(0, begin, end - begin); }); + torch::Tensor indices = batch.m_NonZeroJacobianIndices; + std::lock_guard lock(accumulating); + for (size_t c = 0; c < jacobians.size(); ++c) + { + losses[c]->updateDerivativeInJacobianMode(jacobians[c], indices); + } + }); } template bool ImpactImageToImageMetricv4GetValueAndDerivativeThreader:: - ProcessPoint(const VirtualIndexType & virtualIndex, - const VirtualPointType & virtualPoint, - const FixedImagePointType & mappedFixedPoint, + ProcessPoint(const VirtualIndexType &, + const VirtualPointType &, + const FixedImagePointType &, const FixedImagePixelType &, const FixedImageGradientType &, - const MovingImagePointType & mappedMovingPoint, + const MovingImagePointType &, const MovingImagePixelType &, const MovingImageGradientType &, - MeasureType & metricValueReturn, + MeasureType &, DerivativeType &, - const ThreadIdType threadId) const + const ThreadIdType) const { - // Static mode only: interpolate the (optionally subsampled) per-layer feature vector at the - // mapped fixed/moving point and accumulate the per-layer loss. When a derivative is - // requested, chain the moving feature spatial gradient with the moving-transform Jacobian. - // Value and derivative are reduced and normalized in AfterThreadedExecution. Online mode - // never reaches here: ThreadedExecution() walks its points itself, in batches. - LossPerThreadStruct & loss = this->m_LossThreadStruct[threadId]; + itkExceptionMacro("ProcessPoint is not used: ThreadedExecution evaluates the points in batches."); +} + + +template +void +ImpactImageToImageMetricv4GetValueAndDerivativeThreader:: + EvaluateStaticBatch(typename std::vector::const_iterator first, + typename std::vector::const_iterator last, + const ThreadIdType threadId, + LossPerThreadStruct & loss) const +{ + // Interpolate the per-layer feature vectors at the mapped fixed and moving points and + // accumulate the per-layer losses. When a derivative is requested, chain the moving + // features' spatial gradient with the moving-transform Jacobian. Value and derivative are + // reduced and normalized in AfterThreadedExecution. + const auto count = static_cast(last - first); + const bool computeDerivative = this->GetComputeDerivative(); + constexpr unsigned int movingDimension = ImageToImageMetricv4Type::MovingImageDimension; + const auto & internals = *this->m_ImpactAssociate->m_Internals; - const bool computeDerivative = this->GetComputeDerivative(); + // A map rebuilt through T_k holds the features of M(T_k(p)): read it at x + (T(x) - T_k(x)), + // which is x at the moment of the refresh and moves with the parameters exactly as T(x) does. + // The map built at Initialize() holds the features of M itself and is read at T(x). + std::vector movingReadPoints; + movingReadPoints.reserve(count); + for (auto sample = first; sample != last; ++sample) + { + MovingImagePointType movingReadPoint = sample->m_MovingPoint; + if (internals.m_RefreshTransform.IsNotNull()) + { + const auto anchored = internals.m_RefreshTransform->TransformPoint(sample->m_VirtualPoint); + for (unsigned int d = 0; d < movingDimension; ++d) + { + movingReadPoint[d] = sample->m_VirtualPoint[d] + (sample->m_MovingPoint[d] - anchored[d]); + } + } + movingReadPoints.push_back(movingReadPoint); + } - torch::Tensor transformJacobian; // [1, MovingDim, P] - torch::Tensor nonZeroJacobianIndices; // [1, P] + torch::Tensor transformJacobian; // [B, MovingDim, P] + torch::Tensor nonZeroJacobianIndices; // [B, P] if (computeDerivative) { using JacobianType = typename TImageToImageMetric::JacobianType; JacobianType & jacobian = this->m_GetValueAndDerivativePerThreadVariables[threadId].MovingTransformJacobian; JacobianType & jacobianPositional = this->m_GetValueAndDerivativePerThreadVariables[threadId].MovingTransformJacobianPositional; - // For a dense/global transform this is the full Jacobian and every parameter - // contributes at every point (handled below via an identity index list). - this->m_ImpactAssociate->GetMovingTransform()->ComputeJacobianWithRespectToParametersCachedTemporaries( - virtualPoint, jacobian, jacobianPositional); - - constexpr unsigned int movingDimension = ImageToImageMetricv4Type::MovingImageDimension; - const NumberOfParametersType numberOfLocalParameters = this->GetCachedNumberOfLocalParameters(); - transformJacobian = torch::empty( - { 1, static_cast(movingDimension), static_cast(numberOfLocalParameters) }, torch::kFloat32); - auto jacobianAccessor = transformJacobian.accessor(); - for (unsigned int dimension = 0; dimension < movingDimension; ++dimension) + const auto transform = this->m_ImpactAssociate->GetMovingTransform(); + const auto parameters = static_cast(this->GetCachedNumberOfLocalParameters()); + // A global transform (affine, rigid) has offset 0. A local-support displacement field only + // affects the parameter block of the point's own virtual voxel. + using MovingTransformType = typename TImageToImageMetric::MovingTransformType; + const bool field = + transform->GetTransformCategory() == MovingTransformType::TransformCategoryEnum::DisplacementField; + + transformJacobian = torch::empty({ count, movingDimension, parameters }, torch::kFloat32); + nonZeroJacobianIndices = torch::empty({ count, parameters }, torch::kInt64); + float * jacobianOut = transformJacobian.data_ptr(); + int64_t * indicesOut = nonZeroJacobianIndices.data_ptr(); + for (auto sample = first; sample != last; ++sample) { - for (NumberOfParametersType parameter = 0; parameter < numberOfLocalParameters; ++parameter) + transform->ComputeJacobianWithRespectToParametersCachedTemporaries( + sample->m_VirtualPoint, jacobian, jacobianPositional); + for (unsigned int d = 0; d < movingDimension; ++d) { - jacobianAccessor[0][dimension][parameter] = static_cast(jacobian(dimension, parameter)); + for (int64_t parameter = 0; parameter < parameters; ++parameter) + { + *jacobianOut++ = static_cast(jacobian(d, parameter)); + } + } + const int64_t offset = + field ? static_cast( + this->m_ImpactAssociate->ComputeParameterOffsetFromVirtualIndex(sample->m_VirtualIndex, parameters)) + : 0; + for (int64_t parameter = 0; parameter < parameters; ++parameter) + { + *indicesOut++ = offset + parameter; } } - // Map local parameter indices to global ones. A global transform (affine, rigid) has - // offset 0. A local-support displacement field only affects the parameter block of - // this point's own virtual voxel, at ComputeParameterOffsetFromVirtualIndex. - int64_t parameterOffset = 0; - using MovingTransformType = typename TImageToImageMetric::MovingTransformType; - if (this->m_ImpactAssociate->GetMovingTransform()->GetTransformCategory() == - MovingTransformType::TransformCategoryEnum::DisplacementField) - { - parameterOffset = static_cast( - this->m_ImpactAssociate->ComputeParameterOffsetFromVirtualIndex(virtualIndex, numberOfLocalParameters)); - } - nonZeroJacobianIndices = - (torch::arange(static_cast(numberOfLocalParameters), torch::kInt64) + parameterOffset).unsqueeze(0); } - // A map rebuilt through T_k holds the features of M(T_k(p)): read it at x + (T(x) - T_k(x)), - // which is x at the moment of the refresh and moves with the parameters exactly as T(x) does. - // The map built at Initialize() holds the features of M itself and is read at T(x). - MovingImagePointType movingReadPoint = mappedMovingPoint; - if (const auto & refresh = this->m_ImpactAssociate->m_Internals->m_RefreshTransform; refresh.IsNotNull()) + for (size_t i = 0; i < internals.m_FixedFeaturesMaps.size(); ++i) { - const auto anchored = refresh->TransformPoint(virtualPoint); - for (unsigned int d = 0; d < MovingImagePointType::PointDimension; ++d) + const std::vector & subset = this->m_SubsetsOfFeatures[i]; + const auto channels = static_cast(subset.size()); + const auto & fixedMap = *internals.m_FixedFeaturesMaps[i].m_FeaturesMapInterpolator; + const auto & movingMap = *internals.m_MovingFeaturesMaps[i].m_FeaturesMapInterpolator; + + torch::Tensor fixedFeatures = torch::empty({ count, channels }, torch::kFloat32); + torch::Tensor movingFeatures = torch::empty({ count, channels }, torch::kFloat32); + torch::Tensor movingDerivatives; // d(feature)/d(moving coordinate) : [B, C, MovingDim] + float * fixedOut = fixedFeatures.data_ptr(); + float * movingOut = movingFeatures.data_ptr(); + float * derivativeOut = nullptr; + if (computeDerivative) { - movingReadPoint[d] = virtualPoint[d] + (mappedMovingPoint[d] - anchored[d]); + movingDerivatives = torch::empty({ count, channels, movingDimension }, torch::kFloat32); + derivativeOut = movingDerivatives.data_ptr(); + } + int64_t row = 0; + for (auto sample = first; sample != last; ++sample, ++row) + { + fixedMap.Evaluate(sample->m_FixedPoint, subset, fixedOut + row * channels); + if (computeDerivative) + { + movingMap.EvaluateValueAndDerivative( + movingReadPoints[row], subset, movingOut + row * channels, derivativeOut + row * channels * movingDimension); + } + else + { + movingMap.Evaluate(movingReadPoints[row], subset, movingOut + row * channels); + } } - } - - for (size_t i = 0; i < this->m_ImpactAssociate->m_Internals->m_FixedFeaturesMaps.size(); ++i) - { - const std::vector subsetOfFeatures = - this->GetSubsetOfFeatures(this->m_ImpactAssociate->m_features_indexes[i], - loss.m_randomGenerator, - this->m_ImpactAssociate->GetSubsetFeatures()[i]); - - torch::Tensor fixedFeatures = this->m_ImpactAssociate->m_Internals->m_FixedFeaturesMaps[i] - .m_FeaturesMapInterpolator->Evaluate(mappedFixedPoint, subsetOfFeatures) - .unsqueeze(0); - torch::Tensor movingFeatures = this->m_ImpactAssociate->m_Internals->m_MovingFeaturesMaps[i] - .m_FeaturesMapInterpolator->Evaluate(movingReadPoint, subsetOfFeatures) - .unsqueeze(0); if (computeDerivative) { - // d(feature)/d(moving coordinate) : [1, C, MovingDim] - torch::Tensor movingFeatureDerivative = - this->m_ImpactAssociate->m_Internals->m_MovingFeaturesMaps[i] - .m_FeaturesMapInterpolator->EvaluateDerivative(movingReadPoint, subsetOfFeatures) - .unsqueeze(0); - // chain with d(moving coordinate)/d(parameter) : [1, C, P] - torch::Tensor featureParameterJacobian = torch::bmm(movingFeatureDerivative, transformJacobian); + // chained with d(moving coordinate)/d(parameter) : [B, C, P] + torch::Tensor featureParameterJacobian = torch::bmm(movingDerivatives, transformJacobian); loss.m_losses[i]->updateValueAndDerivativeInStaticMode( fixedFeatures, movingFeatures, featureParameterJacobian, nonZeroJacobianIndices); } @@ -774,10 +914,6 @@ ImpactImageToImageMetricv4GetValueAndDerivativeThreaderupdateValue(fixedFeatures, movingFeatures); } } - ++loss.m_numberOfPixelsCounted; - - metricValueReturn = MeasureType{}; - return true; } } // end namespace itk diff --git a/include/itkInterpolateVectorImageFunction.h b/include/itkInterpolateVectorImageFunction.h index 18ff2ed..fc6bba2 100644 --- a/include/itkInterpolateVectorImageFunction.h +++ b/include/itkInterpolateVectorImageFunction.h @@ -19,11 +19,14 @@ /** * \class InterpolateVectorImageFunction * \ingroup Impact - * \brief Interpolate each channel of a VectorImage with its own B-Spline interpolator. + * \brief Interpolate the channels of a VectorImage with cubic B-splines. * - * ITK's BSplineInterpolateImageFunction does not support VectorImages, so one interpolator is - * built per channel. Evaluation at arbitrary physical points returns the interpolated values or - * spatial derivatives as torch tensors. + * ITK's BSplineInterpolateImageFunction does not support VectorImages. This class computes each + * channel's B-spline coefficients as TInterpolator (a BSplineInterpolateImageFunction of order 3) + * does, and evaluates them as it does, to the bit. But the support, the weights and the boundary + * mirroring are the same for every channel, so they are computed once per point, and the + * coefficients of a voxel are stored channel after channel. Evaluation at arbitrary physical + * points returns the interpolated values or spatial derivatives as torch tensors or into buffers. */ #ifndef itkInterpolateVectorImageFunction_h #define itkInterpolateVectorImageFunction_h @@ -33,6 +36,10 @@ #include #include +#include +#include +#include + namespace itk { @@ -59,7 +66,6 @@ class ITK_TEMPLATE_EXPORT InterpolateVectorImageFunction using ImageType = TImage; using InterpolatorType = TInterpolator; - using InterpolatorPointer = typename InterpolatorType::Pointer; using ImagePointer = typename ImageType::Pointer; using PixelType = typename ImageType::PixelType; using ImagePointType = typename ImageType::PointType; @@ -67,8 +73,6 @@ class ITK_TEMPLATE_EXPORT InterpolateVectorImageFunction /** Dimension underlying input image. */ static constexpr unsigned int ImageDimension = Superclass::ImageDimension; - using CovariantVectorType = itk::CovariantVector; - InterpolateVectorImageFunction() = default; @@ -76,7 +80,7 @@ class ITK_TEMPLATE_EXPORT InterpolateVectorImageFunction using Superclass::SetInputImage; /** - * \brief Initialize one B-Spline interpolator per feature channel of the input VectorImage. + * \brief Compute the cubic B-spline coefficients of every channel of the input VectorImage. * * \param vectorImage The input VectorImage whose channels are interpolated. */ @@ -92,7 +96,7 @@ class ITK_TEMPLATE_EXPORT InterpolateVectorImageFunction * \return A 1D torch::Tensor containing interpolated values for the requested channels. */ torch::Tensor - Evaluate(ImagePointType point, std::vector subsetOfFeatures) const; + Evaluate(const ImagePointType & point, const std::vector & subsetOfFeatures) const; /** * \brief Evaluates the spatial derivative of selected features at a given point. @@ -106,7 +110,25 @@ class ITK_TEMPLATE_EXPORT InterpolateVectorImageFunction * \return A 2D torch::Tensor (Channels × SpatialDimension) with spatial gradients per feature. */ torch::Tensor - EvaluateDerivative(ImagePointType point, std::vector subsetOfFeatures) const; + EvaluateDerivative(const ImagePointType & point, const std::vector & subsetOfFeatures) const; + + /** Keep a copy of these channels' coefficients voxel after voxel, the channels of a voxel side by side: + * evaluating exactly these channels then reads the fewest cache lines. Any other subset still reads the + * coefficients one channel after the other. Not thread-safe: call it before evaluating from threads. */ + void + SetActiveChannels(const std::vector & subsetOfFeatures); + + /** Evaluate() written to `values`, one float per selected channel, without allocating a tensor. */ + void + Evaluate(const ImagePointType & point, const std::vector & subsetOfFeatures, float * values) const; + + /** The selected channels' values to `values` and their spatial derivatives to `derivatives` (ImageDimension + * floats per channel, channel after channel), sharing each channel's B-spline weights between the two. */ + void + EvaluateValueAndDerivative(const ImagePointType & point, + const std::vector & subsetOfFeatures, + float * values, + float * derivatives) const; using OutputType = typename Superclass::OutputType; using ContinuousIndexType = typename Superclass::ContinuousIndexType; @@ -125,8 +147,18 @@ class ITK_TEMPLATE_EXPORT InterpolateVectorImageFunction } private: - ImagePointer m_VectorImage; - std::vector m_Interpolators; + /** The values, and with VDerivative the spatial derivatives, of the selected channels. */ + template + void + EvaluateChannels(const ImagePointType & point, + const std::vector & subsetOfFeatures, + float * values, + float * derivatives) const; + + ImagePointer m_VectorImage; + std::vector m_Coefficients; // one channel after the other, each voxel after voxel (axis 0 fastest) + std::vector m_ActiveChannels; + std::vector m_ActiveCoefficients; // voxel after voxel, the active channels side by side }; } // end namespace itk diff --git a/include/itkInterpolateVectorImageFunction.hxx b/include/itkInterpolateVectorImageFunction.hxx index 8dba9df..06e099a 100644 --- a/include/itkInterpolateVectorImageFunction.hxx +++ b/include/itkInterpolateVectorImageFunction.hxx @@ -28,63 +28,274 @@ template void InterpolateVectorImageFunction::SetInputImage(ImagePointer vectorImage) { - // One scalar image + interpolator per channel. - for (unsigned int i = 0; i < vectorImage->GetVectorLength(); ++i) + // Each channel's coefficients as TInterpolator computes them, one channel after the other: an + // evaluation reads only the channels it compares. + using ChannelImageType = itk::Image; + const unsigned int channels = vectorImage->GetVectorLength(); + const size_t voxels = vectorImage->GetBufferedRegion().GetNumberOfPixels(); + std::vector coefficients(voxels * channels); + for (unsigned int channel = 0; channel < channels; ++channel) { - auto selector = itk::VectorIndexSelectionCastImageFilter>::New(); + auto selector = itk::VectorIndexSelectionCastImageFilter::New(); selector->SetInput(vectorImage); - selector->SetIndex(i); - selector->Update(); - - auto interpolator = TInterpolator::New(); - interpolator->SetInputImage(selector->GetOutput()); - interpolator->SetSplineOrder(3); - this->m_Interpolators.push_back(interpolator); + selector->SetIndex(channel); + auto decomposition = TInterpolator::CoefficientFilter::New(); + decomposition->SetInput(selector->GetOutput()); + decomposition->SetSplineOrder(3); + decomposition->Update(); + const auto * channelCoefficients = decomposition->GetOutput()->GetBufferPointer(); + std::copy(channelCoefficients, channelCoefficients + voxels, coefficients.begin() + channel * voxels); } + this->m_VectorImage = vectorImage; + this->m_Coefficients = std::move(coefficients); + this->m_ActiveChannels.clear(); + this->m_ActiveCoefficients.clear(); } // end SetInputImage /** - * ******************* Evaluate *********************** + * ******************* SetActiveChannels *********************** */ template -torch::Tensor -InterpolateVectorImageFunction::Evaluate(ImagePointType point, - std::vector subsetOfFeatures) const +void +InterpolateVectorImageFunction::SetActiveChannels( + const std::vector & subsetOfFeatures) { - std::vector result; - for (const unsigned int feature : subsetOfFeatures) + if (subsetOfFeatures == this->m_ActiveChannels) { - result.push_back(this->m_Interpolators[feature]->Evaluate(point)); + return; } - return torch::from_blob(result.data(), { static_cast(result.size()) }, torch::kFloat32).clone(); -} // end Evaluate + const size_t voxels = this->m_VectorImage->GetBufferedRegion().GetNumberOfPixels(); + const size_t channels = subsetOfFeatures.size(); + this->m_ActiveCoefficients.resize(voxels * channels); + for (size_t c = 0; c < channels; ++c) + { + const float * channel = this->m_Coefficients.data() + subsetOfFeatures[c] * voxels; + for (size_t voxel = 0; voxel < voxels; ++voxel) + { + this->m_ActiveCoefficients[voxel * channels + c] = channel[voxel]; + } + } + this->m_ActiveChannels = subsetOfFeatures; +} // end SetActiveChannels /** - * ******************* EvaluateDerivative *********************** + * ******************* EvaluateChannels *********************** */ template -torch::Tensor -InterpolateVectorImageFunction::EvaluateDerivative( - ImagePointType point, - std::vector subsetOfFeatures) const +template +void +InterpolateVectorImageFunction::EvaluateChannels( + const ImagePointType & point, + const std::vector & subsetOfFeatures, + float * values, + float * derivatives) const { + // BSplineInterpolateImageFunction, order 3, step by step and in the same order of operations: + // Evaluate(), DetermineRegionOfSupport(), SetInterpolationWeights(), SetDerivativeWeights(), + // ApplyMirrorBoundaryConditions(), then EvaluateValueAndDerivativeAtContinuousIndexInternal(). + constexpr unsigned int Dimension = ImageDimension; + constexpr unsigned int Support = 4; + using CoordRepType = typename TInterpolator::CoordRepType; + const Point coordinatePoint(point); + const auto x = this->m_VectorImage->template TransformPhysicalPointToContinuousIndex(coordinatePoint); + const auto & region = this->m_VectorImage->GetBufferedRegion(); + size_t offsets[Dimension][Support]; + double weights[Dimension][Support]; + double derivativeWeights[Dimension][Support]; + size_t stride = 1; + for (unsigned int n = 0; n < Dimension; ++n) + { + long index[Support]; + index[0] = static_cast(std::floor(static_cast(x[n]))) - 1; + for (unsigned int k = 1; k < Support; ++k) + { + index[k] = index[k - 1] + 1; + } - std::vector derivative(subsetOfFeatures.size() * TImage::ImageDimension, 0.0f); - CovariantVectorType dev; - // Fill the derivative tensor with directional gradients for each selected feature - for (int i = 0; i < subsetOfFeatures.size(); ++i) + double w = x[n] - static_cast(index[1]); + weights[n][3] = (1.0 / 6.0) * w * w * w; + weights[n][0] = (1.0 / 6.0) + 0.5 * w * (w - 1.0) - weights[n][3]; + weights[n][2] = w + weights[n][0] - 2.0 * weights[n][3]; + weights[n][1] = 1.0 - weights[n][0] - weights[n][2] - weights[n][3]; + if constexpr (VDerivative) + { + w = x[n] + .5 - static_cast(index[2]); + const double w2 = 0.75 - w * w; + const double w3 = 0.5 * (w - w2 + 1.0); + const double w1 = 1.0 - w2 - w3; + derivativeWeights[n][0] = 0.0 - w1; + derivativeWeights[n][1] = w1 - w2; + derivativeWeights[n][2] = w2 - w3; + derivativeWeights[n][3] = w3; + } + + const long start = region.GetIndex(n); + const long end = start + static_cast(region.GetSize(n)) - 1; + for (unsigned int k = 0; k < Support; ++k) + { + if (region.GetSize(n) == 1) + { + index[k] = start; + } + else + { + if (index[k] < start) + { + index[k] = start + (start - index[k]); + } + if (index[k] >= end) + { + index[k] = end - (index[k] - end); + } + } + offsets[n][k] = static_cast(index[k] - start) * stride; + } + stride *= region.GetSize(n); + } + + // Every point of the support, axis 0 fastest, accumulated in double; the channels a block at a + // time, so that the accumulators stay on the stack. + constexpr size_t Block = 64; + const size_t voxels = stride; // the coefficients of one channel + const bool active = subsetOfFeatures == this->m_ActiveChannels; + const auto & spacing = this->m_VectorImage->GetSpacing(); + const auto & direction = this->m_VectorImage->GetDirection(); + for (size_t first = 0; first < subsetOfFeatures.size(); first += Block) { - dev = this->m_Interpolators[subsetOfFeatures[i]]->EvaluateDerivative(point); - for (unsigned int it = 0; it < TImage::ImageDimension; ++it) + const size_t channels = std::min(Block, subsetOfFeatures.size() - first); + const auto * channel = subsetOfFeatures.data() + first; + double value[Block] = {}; + double derivative[Dimension][Block] = {}; + unsigned int k[Dimension] = {}; + for (unsigned int p = 0; p < (1u << (2 * Dimension)); ++p) + { + double w = weights[0][k[0]]; + size_t offset = offsets[0][k[0]]; + for (unsigned int n = 1; n < Dimension; ++n) + { + w *= weights[n][k[n]]; + offset += offsets[n][k[n]]; + } + double derivativeW[Dimension]; + if constexpr (VDerivative) + { + for (unsigned int n = 0; n < Dimension; ++n) + { + derivativeW[n] = 1.0; + for (unsigned int n1 = 0; n1 < Dimension; ++n1) + { + derivativeW[n] *= n1 == n ? derivativeWeights[n1][k[n1]] : weights[n1][k[n1]]; + } + } + } + auto accumulate = [&](size_t c, double coefficient) { + value[c] += w * coefficient; + if constexpr (VDerivative) + { + for (unsigned int n = 0; n < Dimension; ++n) + { + derivative[n][c] += derivativeW[n] * coefficient; + } + } + }; + if (active) + { + const float * coefficients = this->m_ActiveCoefficients.data() + offset * subsetOfFeatures.size() + first; + for (size_t c = 0; c < channels; ++c) + { + accumulate(c, coefficients[c]); + } + } + else + { + const float * coefficients = this->m_Coefficients.data() + offset; + for (size_t c = 0; c < channels; ++c) + { + accumulate(c, coefficients[channel[c] * voxels]); + } + } + for (unsigned int n = 0; n < Dimension && ++k[n] == Support; ++n) + { + k[n] = 0; + } + } + + for (size_t c = 0; c < channels; ++c) { - derivative[i * TImage::ImageDimension + it] = static_cast(dev[it]); + values[first + c] = static_cast(value[c]); + if constexpr (VDerivative) + { + double local[Dimension]; + for (unsigned int n = 0; n < Dimension; ++n) + { + local[n] = derivative[n][c] / spacing[n]; + } + for (unsigned int i = 0; i < Dimension; ++i) + { + double sum = 0.0; + for (unsigned int j = 0; j < Dimension; ++j) + { + sum += direction[i][j] * local[j]; + } + derivatives[(first + c) * Dimension + i] = static_cast(sum); + } + } } } - return torch::from_blob(derivative.data(), - { static_cast(subsetOfFeatures.size()), TImage::ImageDimension }, - torch::kFloat32) - .clone(); +} // end EvaluateChannels + +/** + * ******************* Evaluate *********************** + */ +template +torch::Tensor +InterpolateVectorImageFunction::Evaluate( + const ImagePointType & point, + const std::vector & subsetOfFeatures) const +{ + torch::Tensor values = torch::empty({ static_cast(subsetOfFeatures.size()) }, torch::kFloat32); + this->Evaluate(point, subsetOfFeatures, values.data_ptr()); + return values; +} // end Evaluate + +template +void +InterpolateVectorImageFunction::Evaluate(const ImagePointType & point, + const std::vector & subsetOfFeatures, + float * values) const +{ + this->template EvaluateChannels(point, subsetOfFeatures, values, nullptr); +} // end Evaluate + +/** + * ******************* EvaluateValueAndDerivative *********************** + */ +template +void +InterpolateVectorImageFunction::EvaluateValueAndDerivative( + const ImagePointType & point, + const std::vector & subsetOfFeatures, + float * values, + float * derivatives) const +{ + this->template EvaluateChannels(point, subsetOfFeatures, values, derivatives); +} // end EvaluateValueAndDerivative + +/** + * ******************* EvaluateDerivative *********************** + */ +template +torch::Tensor +InterpolateVectorImageFunction::EvaluateDerivative( + const ImagePointType & point, + const std::vector & subsetOfFeatures) const +{ + const auto channels = static_cast(subsetOfFeatures.size()); + std::vector values(channels); + torch::Tensor derivatives = torch::empty({ channels, TImage::ImageDimension }, torch::kFloat32); + this->EvaluateValueAndDerivative(point, subsetOfFeatures, values.data(), derivatives.data_ptr()); + return derivatives; } // end EvaluateDerivative } // namespace itk #endif // end #ifndef itkInterpolateVectorImageFunction_hxx diff --git a/test/itkImpactBackendGTest.cxx b/test/itkImpactBackendGTest.cxx index a122eee..f2c1cd6 100644 --- a/test/itkImpactBackendGTest.cxx +++ b/test/itkImpactBackendGTest.cxx @@ -969,9 +969,10 @@ TEST(ImpactMetric, JacobianResultIsIndependentOfBatchSize) metric->SetModelsConfiguration(configs); metric->SetDistance({ std::string(lossName) }); metric->SetLayersWeight({ 1.f }); - // All four channels: the feature subset is then drawn without consuming the generator, so a - // different number of batches cannot shift the random sequence and confound the comparison. - metric->SetSubsetFeatures({ 4 }); + // Two of the four channels: every point of an evaluation compares the same two, whatever the + // batches. NCC accumulates its statistics channel by channel, so a subset drawn per batch + // mixed channels within one sum. + metric->SetSubsetFeatures({ 2 }); metric->SetPCA({ 0 }); metric->SetMode("Jacobian"); metric->SetSeed(1); @@ -1030,6 +1031,141 @@ TEST(ImpactMetric, JacobianResultIsIndependentOfBatchSize) } } +// On a GPU the online mode runs every point through the models in batches as large as the device +// holds: the value and the derivative have to be the CPU's, up to the TF32 arithmetic of CUDA +// convolutions. Skipped without a CUDA device. +TEST(ImpactMetric, JacobianOnTheGpuMatchesTheCpu) +{ + if (!torch::cuda::is_available()) + { + GTEST_SKIP() << "no CUDA device"; + } + using VirtualImageType = itk::Image; + using MetricType = itk::ImpactImageToImageMetricv4; + + auto fixed = MakeRampImage(8, 0); + auto moving = MakeRampImage(8, 1); + + auto evaluate = [&](const char * lossName, const char * device, MetricType::DerivativeType & derivative) { + auto metric = MetricType::New(); + std::vector configs; + configs.emplace_back(ToyModelPath(), + 3, + 1, + std::vector{ 3, 3, 3 }, + std::vector{ 1.f, 1.f, 1.f }, + std::vector{ 0, 0, 0 }, + std::vector{ true, false }, + false); + metric->SetModelsConfiguration(configs); + metric->SetDistance({ std::string(lossName) }); + metric->SetLayersWeight({ 1.f }); + metric->SetSubsetFeatures({ 2 }); + metric->SetPCA({ 0 }); + metric->SetMode("Jacobian"); + metric->SetSeed(1); + metric->SetDevice(device); + + auto movingInterp = itk::BSplineInterpolateImageFunction::New(); + movingInterp->SetSplineOrder(3); + metric->SetMovingInterpolator(movingInterp); + auto transform = itk::TranslationTransform::New(); + transform->SetIdentity(); + metric->SetFixedImage(fixed); + metric->SetMovingImage(moving); + metric->SetFixedTransform(itk::IdentityTransform::New()); + metric->SetMovingTransform(transform); + metric->Initialize(); + + MetricType::MeasureType value; + metric->GetValueAndDerivative(value, derivative); + return static_cast(value); + }; + + for (const char * lossName : { "L2", "NCC" }) + { + SCOPED_TRACE(lossName); + MetricType::DerivativeType cpu, gpu; + const double cpuValue = evaluate(lossName, "cpu", cpu); + const double gpuValue = evaluate(lossName, "cuda", gpu); + EXPECT_NEAR(gpuValue, cpuValue, 1e-3 * std::abs(cpuValue)); + ASSERT_EQ(gpu.GetSize(), cpu.GetSize()); + double difference = 0.0, norm = 0.0; + for (unsigned int j = 0; j < cpu.GetSize(); ++j) + { + difference += (gpu[j] - cpu[j]) * (gpu[j] - cpu[j]); + norm += cpu[j] * cpu[j]; + } + ASSERT_GT(norm, 1e-12); + EXPECT_LT(std::sqrt(difference / norm), 1e-2); + } +} + +// The Static analogue: the points are split over the work units, and a feature subset drawn per +// point from each work unit's own generator gave every point a subset that depended on its rank in +// that split. All points of an evaluation now compare the same channels. +TEST(ImpactMetric, StaticResultIsIndependentOfTheNumberOfWorkUnits) +{ + using VirtualImageType = itk::Image; + using MetricType = itk::ImpactImageToImageMetricv4; + + auto fixed = MakeRampImage(8, 0); + auto moving = MakeRampImage(8, 1); + + auto evaluate = [&](const char * lossName, unsigned int workUnits, MetricType::DerivativeType & derivative) { + auto metric = MetricType::New(); + std::vector configs; + configs.emplace_back(ToyModelPath(), + 3, + 1, + std::vector{ 0, 0, 0 }, + std::vector{ 1.f, 1.f, 1.f }, + std::vector{ 2, 2, 2 }, + std::vector{ true, false }, + false); + metric->SetModelsConfiguration(configs); + metric->SetDistance({ std::string(lossName) }); + metric->SetLayersWeight({ 1.f }); + metric->SetSubsetFeatures({ 2 }); + metric->SetPCA({ 0 }); + metric->SetMode("Static"); + metric->SetSeed(1); + metric->SetFeaturesMapUpdateInterval(-1); + metric->SetDevice("cpu"); + + auto affine = itk::AffineTransform::New(); + affine->SetIdentity(); + + metric->SetFixedImage(fixed); + metric->SetMovingImage(moving); + metric->SetFixedTransform(itk::IdentityTransform::New()); + metric->SetMovingTransform(affine); + metric->SetUseFixedImageGradientFilter(false); + metric->SetUseMovingImageGradientFilter(false); + metric->SetMaximumNumberOfWorkUnits(workUnits); + metric->Initialize(); + + MetricType::MeasureType value; + metric->GetValueAndDerivative(value, derivative); + return static_cast(value); + }; + + // Relative bars, as for the batches: the work units only regroup the float sums. + for (const char * lossName : { "L2", "NCC" }) + { + SCOPED_TRACE(lossName); + MetricType::DerivativeType one, four; + const double oneValue = evaluate(lossName, 1, one); + const double fourValue = evaluate(lossName, 4, four); + + EXPECT_NEAR(fourValue, oneValue, 1e-6 * std::abs(oneValue)); + for (unsigned int j = 0; j < one.GetSize(); ++j) + { + EXPECT_NEAR(four[j], one[j], 1e-4 * std::abs(one[j]) + 1e-12) << "parameter " << j; + } + } +} + // LayersMask is what asks for several depths of the same network to be compared at once -- an // early layer answers for texture, a late one for structure. Each kept layer therefore gets its // own feature map, its own loss and its own weight, and the value is their weighted sum. The @@ -2403,6 +2539,32 @@ TEST(ImpactLoss, LocalNCCComparesWindowsOfDenseMaps) EXPECT_THROW(loss->forwardValue(points, points), std::runtime_error); } +// A loss sums its value over the batch of points it is handed. Summed in float32, the value lost +// digits with the size of the batch, so a work unit evaluated in one batch did not give the value +// of the same points fed one by one, and finite differences of it were mostly rounding. +TEST(ImpactLoss, ValueOfABatchIsTheSumOfItsPoints) +{ + torch::manual_seed(7); + const int64_t N = 20000, C = 4; + torch::Tensor fixed = torch::rand({ N, C }) * 10; // float32, as the metric hands them over + torch::Tensor moving = torch::rand({ N, C }) * 10; + for (const std::string name : { "L1", "L2", "Cosine", "L1Cosine", "Dice" }) + { + SCOPED_TRACE(name); + auto batch = itk::Impact::LossFactory::Instance().Create(name); + batch->updateValue(fixed, moving); + auto pointwise = itk::Impact::LossFactory::Instance().Create(name); + for (int64_t i = 0; i < N; ++i) + { + torch::Tensor fixedPoint = fixed.narrow(0, i, 1); + torch::Tensor movingPoint = moving.narrow(0, i, 1); + pointwise->updateValue(fixedPoint, movingPoint); + } + const double expected = pointwise->GetValue(static_cast(N)); + EXPECT_NEAR(batch->GetValue(static_cast(N)), expected, 1e-9 * std::abs(expected)); + } +} + TEST(ImpactLoss, ForwardValueGradientFiniteDifference) { torch::manual_seed(7); @@ -4422,6 +4584,49 @@ TEST(ImpactConvexAdam, CoarseTakesPCAAndASeededChannelSubset) EXPECT_LT(error[d], 1.5) << "PCA to 2 components, axis " << d; } +// The coarse and fine stages draw from generators of their own: running them leaves torch's global +// generator, which a Python caller shares, where it was. Both used to call torch::manual_seed(Seed). +TEST(ImpactConvexAdam, StagesLeaveTorchsGlobalGeneratorAlone) +{ + ImageType::SpacingType spacing; + spacing.Fill(1.0); + ImageType::DirectionType identity; + identity.SetIdentity(); + auto fixed = MakeTorchAdamPattern(24, 0, 0, 0, spacing, identity); + auto moving = MakeTorchAdamPattern(24, 3.0, -2.5, 1.5, spacing, identity); + const itk::ImpactModelConfiguration model( + ToyModelPath(), 3, 1, { 0, 0, 0 }, { 1.f, 1.f, 1.f }, { 0, 0, 0 }, { true, false }, false); + + auto coarse = itk::ImpactCoarseRegistration::New(); + coarse->SetFixedImage(fixed); + coarse->SetMovingImage(moving); + coarse->AddModelConfiguration(model); + coarse->SetSubsetFeatures({ 2 }); // so that the stage draws + coarse->SetSeed(7); + coarse->SetGridSpacing(2); + coarse->SetDisplacementHalfWidth(3); + + auto fine = TorchAdamFilterType::New(); + fine->SetFixedImage(fixed); + fine->SetMovingImage(moving); + fine->AddModelConfiguration(model); + fine->SetSubsetFeatures({ 2 }); + fine->SetSamplingPercentage(0.3); + fine->SetNumberOfIterations(1); + fine->SetSeed(3); + + for (itk::ProcessObject * stage : + { static_cast(coarse), static_cast(fine) }) + { + SCOPED_TRACE(stage->GetNameOfClass()); + torch::manual_seed(123); + const torch::Tensor expected = torch::rand({ 4 }); + torch::manual_seed(123); + stage->Update(); + EXPECT_TRUE(torch::equal(torch::rand({ 4 }), expected)) << "the stage moved torch's global generator"; + } +} + namespace { // NCC of two same-grid images over the interior (margin away from borders). diff --git a/test/itkInterpolateVectorImageFunctionTest.cxx b/test/itkInterpolateVectorImageFunctionTest.cxx index fef54c9..03fe4ec 100644 --- a/test/itkInterpolateVectorImageFunctionTest.cxx +++ b/test/itkInterpolateVectorImageFunctionTest.cxx @@ -16,60 +16,137 @@ * *=========================================================================*/ -// Self-contained: build a small synthetic VectorImage and interpolate a subset of -// its channels at a physical point. No external data required. - +// Self-contained: build a small synthetic VectorImage and check that interpolating its channels +// together gives, to the bit, what ITK's BSplineInterpolateImageFunction gives channel by channel: +// values and derivatives, on a region that does not start at zero, with anisotropic spacing and a +// rotated direction, at points up to the borders, for more channels than one block of 64, and for +// the active channels as for any other subset. #include "itkInterpolateVectorImageFunction.h" #include "itkBSplineInterpolateImageFunction.h" #include "itkVectorImage.h" #include "itkImage.h" #include "itkImageRegionIteratorWithIndex.h" - +#include "itkVectorIndexSelectionCastImageFilter.h" #include +#include #include +#include +#include #include int itkInterpolateVectorImageFunctionTest(int, char *[]) { constexpr unsigned int Dimension = 3; + constexpr unsigned int Channels = 70; using PixelType = float; + using ChannelImageType = itk::Image; using FeaturesImageType = itk::VectorImage; - using InterpolatorType = itk::BSplineInterpolateImageFunction, float, float>; + using InterpolatorType = itk::BSplineInterpolateImageFunction; using VectorInterpolatorType = itk::InterpolateVectorImageFunction; - auto image = FeaturesImageType::New(); - FeaturesImageType::SizeType size; - size.Fill(8); + auto image = FeaturesImageType::New(); FeaturesImageType::RegionType region; - region.SetSize(size); + region.SetIndex({ { 2, -3, 1 } }); + region.SetSize({ { 9, 10, 11 } }); image->SetRegions(region); - image->SetNumberOfComponentsPerPixel(3); + image->SetNumberOfComponentsPerPixel(Channels); + FeaturesImageType::SpacingType spacing; + spacing[0] = 0.7; + spacing[1] = 1.3; + spacing[2] = 2.1; + image->SetSpacing(spacing); + FeaturesImageType::PointType origin; + origin[0] = -4.0; + origin[1] = 2.5; + origin[2] = 10.0; + image->SetOrigin(origin); + FeaturesImageType::DirectionType direction; + const double angle = 0.4; + direction.SetIdentity(); + direction[0][0] = std::cos(angle); + direction[0][1] = -std::sin(angle); + direction[1][0] = std::sin(angle); + direction[1][1] = std::cos(angle); + image->SetDirection(direction); image->Allocate(); - itk::ImageRegionIteratorWithIndex it(image, region); - for (it.GoToBegin(); !it.IsAtEnd(); ++it) + std::mt19937 generator(3); + std::uniform_real_distribution uniform(-1.f, 1.f); + for (size_t i = 0; i < image->GetPixelContainer()->Size(); ++i) { - const auto i = it.GetIndex(); - itk::VariableLengthVector px(3); - px[0] = static_cast(i[0]); - px[1] = static_cast(i[1]); - px[2] = static_cast(i[2]); - it.Set(px); + image->GetBufferPointer()[i] = uniform(generator); } auto interpolator = VectorInterpolatorType::New(); interpolator->SetInputImage(image); + std::vector references; + for (unsigned int channel = 0; channel < Channels; ++channel) + { + auto selector = itk::VectorIndexSelectionCastImageFilter::New(); + selector->SetInput(image); + selector->SetIndex(channel); + selector->Update(); + references.push_back(InterpolatorType::New()); + references.back()->SetInputImage(selector->GetOutput()); + } - FeaturesImageType::PointType point; - point[0] = 3.5; - point[1] = 3.5; - point[2] = 3.5; - const torch::Tensor values = interpolator->Evaluate(point, std::vector{ 0, 1, 2 }); + std::vector subset(Channels); + std::iota(subset.begin(), subset.end(), 0u); + std::shuffle(subset.begin(), subset.end(), generator); + interpolator->SetActiveChannels(subset); + const std::vector other(subset.rbegin(), subset.rbegin() + 20); // not the active channels - std::cout << "interpolated values: " << values << std::endl; - if (values.numel() != 3) + std::vector values(Channels), derivatives(Channels * Dimension), alone(Channels); + std::vector otherValues(other.size()), otherDerivatives(other.size() * Dimension); + unsigned int mismatches = 0; + for (int pointNumber = 0; pointNumber < 500; ++pointNumber) + { + itk::ContinuousIndex index; + for (unsigned int d = 0; d < Dimension; ++d) + { + const double first = region.GetIndex(d) - 0.5; + index[d] = first + (region.GetSize(d)) * (0.5 + 0.5 * uniform(generator)); + } + FeaturesImageType::PointType point; + image->TransformContinuousIndexToPhysicalPoint(index, point); + + interpolator->EvaluateValueAndDerivative(point, subset, values.data(), derivatives.data()); + interpolator->Evaluate(point, subset, alone.data()); + const torch::Tensor tensor = interpolator->Evaluate(point, subset); + const torch::Tensor tensorDerivative = interpolator->EvaluateDerivative(point, subset); + interpolator->EvaluateValueAndDerivative(point, other, otherValues.data(), otherDerivatives.data()); + for (size_t c = 0; c < other.size(); ++c) + { + const auto & reference = references[other[c]]; + const auto derivative = reference->EvaluateDerivative(point); + bool same = otherValues[c] == static_cast(reference->Evaluate(point)); + for (unsigned int d = 0; d < Dimension; ++d) + { + same = same && otherDerivatives[c * Dimension + d] == static_cast(derivative[d]); + } + mismatches += same ? 0 : 1; + } + for (unsigned int c = 0; c < Channels; ++c) + { + const auto & reference = references[subset[c]]; + const auto value = static_cast(reference->Evaluate(point)); + const auto derivative = reference->EvaluateDerivative(point); + bool same = values[c] == value && alone[c] == value && tensor[c].item() == value; + for (unsigned int d = 0; d < Dimension; ++d) + { + same = same && derivatives[c * Dimension + d] == static_cast(derivative[d]) && + tensorDerivative[c][d].item() == static_cast(derivative[d]); + } + if (!same && ++mismatches <= 5) + { + std::cerr << "channel " << subset[c] << " at " << index << ": " << values[c] << " vs " << value << ", " + << derivatives[c * Dimension] << " vs " << derivative[0] << std::endl; + } + } + } + if (mismatches > 0) { - std::cerr << "InterpolateVectorImageFunction returned an unexpected number of values" << std::endl; + std::cerr << mismatches << " channel evaluations differ from BSplineInterpolateImageFunction" << std::endl; return EXIT_FAILURE; } return EXIT_SUCCESS;