BUG: Draw the feature subsets once per evaluation, and speed up the v4 metric

include/ImpactLoss.h

@@ -314,15 +314,15 @@ 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<double>();
+    this->m_value += (fixedOutput - movingOutput).abs().mean(1).sum(torch::kDouble).item<double>();
   }
 
   torch::Tensor
   updateValueAndGetGradientModulator(torch::Tensor & fixedOutput, torch::Tensor & movingOutput) override
   {
     this->initialize(fixedOutput);
     torch::Tensor diffOutput = fixedOutput - movingOutput;
-    this->m_value += diffOutput.abs().mean(1).sum().item<double>();
+    this->m_value += diffOutput.abs().mean(1).sum(torch::kDouble).item<double>();
     return -torch::sign(diffOutput) / fixedOutput.size(1);
   }
 
@@ -359,15 +359,15 @@ 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<double>();
+    this->m_value += (fixedOutput - movingOutput).pow(2).mean(1).sum(torch::kDouble).item<double>();
   }
 
   torch::Tensor
   updateValueAndGetGradientModulator(torch::Tensor & fixedOutput, torch::Tensor & movingOutput) override
   {
     this->initialize(fixedOutput);
     torch::Tensor diffOutput = fixedOutput - movingOutput;
-    this->m_value += diffOutput.pow(2).mean(1).sum().item<double>();
+    this->m_value += diffOutput.pow(2).mean(1).sum(torch::kDouble).item<double>();
     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<double>();
+    this->m_value += (1.0 - dice).sum(torch::kDouble).item<double>();
   }
 
   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<double>();
+    this->m_value += (1.0 - dice).sum(torch::kDouble).item<double>();
 
     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<double>();
+    this->m_value += (1.0 - (cosine.unsqueeze(-1) * expL1).mean(1)).sum(torch::kDouble).item<double>();
   }
 
   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<double>();
+    this->m_value += (1.0 - (cosine.unsqueeze(-1) * expL1).mean(1)).sum(torch::kDouble).item<double>();
     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<double>();
+    this->m_value += (1.0 - dot_product / (norm_fixed * norm_moving)).sum(torch::kDouble).item<double>();
   }
 
   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<double>();
+    this->m_value += (1.0 - dot_product / v).sum(torch::kDouble).item<double>();
     // 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) -

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<int64_t> & patchSize = configuration.GetPatchSize();
@@ -160,35 +166,41 @@ 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<double>(static_cast<int64_t>(freeBytes) + std::max<int64_t>(cached, 0)) * safety;
-  return FitBatchToDevice(
-    std::max<int64_t>(static_cast<int64_t>(usable / static_cast<double>(perPatch)), 1), usable, [&](int64_t batch) {
-      try
-      {
-        return peakOf(batch);
-      }
-      catch (...)
-      {
-        c10::cuda::CUDACachingAllocator::emptyCache();
-        throw;
-      }
-    });
+  int64_t       estimate = std::max<int64_t>(static_cast<int64_t>(usable / static_cast<double>(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<ImpactModelConfiguration> & configurations,
                    const torch::Device &                   device,
                    int64_t                                 requested = 0)
 {
   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);
   }

include/itkImpactCoarseRegistration.hxx

@@ -28,6 +28,7 @@
 #include <itkIdentityTransform.h>
 
 #include <torch/torch.h>
+#include <ATen/CPUGeneratorImpl.h>
 
 #include <algorithm>
 #include <cmath>
@@ -325,7 +326,10 @@ ImpactCoarseRegistration<TFixedImage, TMovingImage>::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<at::CPUGeneratorImpl>(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<TFixedImage, TMovingImage>::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();
         }

include/itkImpactFineRegistration.hxx

@@ -31,6 +31,7 @@
 #include <itkIdentityTransform.h>
 
 #include <torch/torch.h>
+#include <ATen/CPUGeneratorImpl.h>
 
 #include <algorithm>
 #include <array>
@@ -221,7 +222,10 @@ ImpactFineRegistration<TFixedImage, TMovingImage>::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<at::CPUGeneratorImpl>(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<TFixedImage, TMovingImage>::GenerateData()
       const int64_t kept = l < m_SubsetFeatures.size() ? static_cast<int64_t>(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<TFixedImage, TMovingImage>::GenerateData()
       }
       count = std::max<int64_t>(1, std::llround(m_SamplingPercentage * static_cast<double>(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<int64_t> pointShape(ImageDimension + 2, 1); // grid_sample's {1, [1,] 1, n, N}
@@ -1036,7 +1039,7 @@ ImpactFineRegistration<TFixedImage, TMovingImage>::GenerateData()
           voxels *= extent;
         }
         const int64_t count = std::max<int64_t>(1, std::llround(m_SamplingPercentage * static_cast<double>(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<int64_t>(ImageDimension) }, torch::kLong);
         int64_t             stride = 1;
         for (unsigned int d = 0; d < ImageDimension; ++d) // ITK index, x fastest in the flat voxel number

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<std::vector<unsigned int>> m_features_indexes;
 };

include/itkImpactImageToImageMetricv4.hxx

@@ -95,6 +95,10 @@ ImpactImageToImageMetricv4<TFixedImage, TMovingImage, TVirtualImage, TInternalCo
     ImpactImageToImageMetricv4GetValueAndDerivativeThreader<ThreadedIndexedContainerPartitioner, Superclass, Self>;
   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 <typename TFixedImage,
@@ -489,10 +493,13 @@ ImpactImageToImageMetricv4<TFixedImage, TMovingImage, TVirtualImage, TInternalCo
       }
     }
 
-    // Device budget per model (itkImpactBatchBudget.h); its backward pass refuses the GIL.
+    // Device budget per model (itkImpactBatchBudget.h), up to the largest batch an evaluation can
+    // run: BatchSize, or all the points; its backward pass refuses the GIL.
+    const auto largest =
+      static_cast<int64_t>(this->m_BatchSize > 0 ? this->m_BatchSize : this->GetNumberOfDomainPoints());
     Impact::PythonGilReleaseGuard noGil;
-    Impact::ConfigureBatchSize(this->m_MovingModelsConfiguration, device, static_cast<int64_t>(this->m_BatchSize));
-    Impact::ConfigureBatchSize(this->m_FixedModelsConfiguration, device, static_cast<int64_t>(this->m_BatchSize));
+    Impact::ConfigureBatchSize(this->m_MovingModelsConfiguration, device, largest);
+    Impact::ConfigureBatchSize(this->m_FixedModelsConfiguration, device, largest);
   }
   else
   {