diffct 2.3.0__tar.gz → 2.4.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {diffct-2.3.0 → diffct-2.4.0}/CHANGELOG.md +26 -0
- {diffct-2.3.0 → diffct-2.4.0}/CITATION.cff +5 -5
- {diffct-2.3.0 → diffct-2.4.0}/PKG-INFO +24 -18
- {diffct-2.3.0 → diffct-2.4.0}/README.md +23 -17
- {diffct-2.3.0 → diffct-2.4.0}/diffct/__init__.py +5 -1
- {diffct-2.3.0 → diffct-2.4.0}/diffct/_solver_state.py +42 -25
- {diffct-2.3.0 → diffct-2.4.0}/diffct/backends.py +140 -50
- {diffct-2.3.0 → diffct-2.4.0}/diffct/chunking.py +34 -0
- {diffct-2.3.0 → diffct-2.4.0}/diffct/distribution.py +67 -7
- {diffct-2.3.0 → diffct-2.4.0}/diffct/geometry.py +93 -17
- {diffct-2.3.0 → diffct-2.4.0}/diffct/kernels/cone_beam.py +21 -8
- diffct-2.4.0/diffct/kernels/surface_geometry.py +130 -0
- {diffct-2.3.0 → diffct-2.4.0}/diffct/operators.py +765 -343
- {diffct-2.3.0 → diffct-2.4.0}/diffct/projectors.py +61 -30
- {diffct-2.3.0 → diffct-2.4.0}/diffct/solvers.py +18 -11
- diffct-2.4.0/diffct/surface_gradients.py +117 -0
- diffct-2.4.0/diffct/surfaces.py +146 -0
- diffct-2.4.0/docs/video/README.md +244 -0
- {diffct-2.3.0 → diffct-2.4.0}/pyproject.toml +1 -1
- diffct-2.4.0/tests/test_call_trajectory.py +252 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_chunked_projector.py +40 -25
- diffct-2.4.0/tests/test_cone_backprojection_transpose.py +189 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_distributed_training_transforms.py +5 -0
- diffct-2.4.0/tests/test_geometry_conversions.py +178 -0
- diffct-2.4.0/tests/test_native_volume_layout.py +328 -0
- diffct-2.4.0/tests/test_native_volume_layout_distributed.py +79 -0
- diffct-2.4.0/tests/test_parallel3d.py +376 -0
- diffct-2.4.0/tests/test_parallel3d_surface_oracle.py +149 -0
- diffct-2.4.0/tests/test_plan_memory_bounds.py +252 -0
- diffct-2.4.0/tests/test_surface_collective_failures.py +521 -0
- diffct-2.4.0/tests/test_surface_composition_kernel.py +358 -0
- diffct-2.4.0/tests/test_surface_compute_shards.py +564 -0
- diffct-2.4.0/tests/test_surface_geometry_cudasim.py +514 -0
- diffct-2.4.0/tests/test_surface_gradient_replay.py +226 -0
- diffct-2.4.0/tests/test_surface_native_regressions.py +281 -0
- diffct-2.4.0/tests/test_surface_pipeline_contracts.py +388 -0
- diffct-2.4.0/tests/test_surface_trajectory_transform.py +142 -0
- diffct-2.4.0/tests/test_trajectory_admission_cost.py +257 -0
- diffct-2.4.0/tests/test_trajectory_surface_integration.py +615 -0
- diffct-2.3.0/diffct/surfaces.py +0 -103
- diffct-2.3.0/docs/video/README.md +0 -119
- {diffct-2.3.0 → diffct-2.4.0}/.gitignore +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/LICENSE +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/NOTICE +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/diffct/_solver_tv.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/diffct/analytical.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/diffct/analytical_execution.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/diffct/constants.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/diffct/differentiable.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/diffct/kernels/__init__.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/diffct/kernels/analytical_transpose.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/diffct/kernels/fan_beam.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/diffct/kernels/geometry_vjp.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/diffct/kernels/parallel_beam.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/diffct/pipeline.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/diffct/reconstruction.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/diffct/storage.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/diffct/streaming.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/diffct/utils.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/examples/README.md +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/examples/_block_cgls.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/examples/_common.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/examples/data/NOTICE +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/examples/data/preprocess_walnut.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/examples/walnut_reconstruction.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/formal/README.md +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/pytest.ini +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/__init__.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/_phantoms.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/benchmarks/README.md +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/benchmarks/__init__.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/benchmarks/_common.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/benchmarks/test_bench_cone.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/benchmarks/test_bench_fan.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/benchmarks/test_bench_parallel.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/distributed_projector_check.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_adjoint_inner_product.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_allocator_capacity.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_analytical_autograd.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_analytical_distribution.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_analytical_regressions.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_analytical_weight_distribution.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_angular_low_precision.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_axis_aligned_chords.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_batch_execution_edges.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_batched_surfaces.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_block_cgls.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_block_storage.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_cgls_state_admission.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_completed_transfer_reuse.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_cone_projector_autograd.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_cpu_install.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_cuda_exit_deallocations.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_cuda_sink_lifetimes.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_cuda_smoke.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_detector_surface_dtypes.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_detector_surface_mapping.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_detector_surface_units.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_detector_surfaces.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_example_psnr.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_execution_buffers.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_execution_features.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_execution_schedules.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_fbp_fan_accuracy.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_fbp_fan_offsets.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_fbp_parallel_accuracy.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_fdk_cone_accuracy.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_fdk_cone_offsets.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_geometry_autograd.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_geometry_gradients.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_gradcheck.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_lazy_native_observers.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_multi_gpu_exit.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_native_surface_transforms.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_output_pipeline_overlap.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_partition_lifecycle.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_portable_distributed_geometry.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_portable_endpoint_precision.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_portable_execution_bounds.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_portable_numeric_regressions.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_portable_surface_admission.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_projector_api.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_projector_degenerate_geometry.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_projector_geometry_devices.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_ramp_filter_windows.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_reconstruction_api.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_reference_backend.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_release_2_0_1.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_review_230.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_review_regressions.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_solver_admission_regressions.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_solver_api.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_solver_checkpoint_edges.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_solver_collective_failures.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_spatial_partition.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_training_transforms.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_transfer_pipeline.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_transfer_retirement_lifetimes.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_tv_zero_dependency.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_two_far_endpoints.py +0 -0
- {diffct-2.3.0 → diffct-2.4.0}/tests/test_weights.py +0 -0
|
@@ -5,6 +5,32 @@ All notable changes to this project are documented in this file.
|
|
|
5
5
|
The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
|
|
6
6
|
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
|
|
7
7
|
|
|
8
|
+
## [Unreleased]
|
|
9
|
+
|
|
10
|
+
## [2.4.0] - 2026-10-20
|
|
11
|
+
|
|
12
|
+
### Added
|
|
13
|
+
|
|
14
|
+
- Per-call `trajectory` overrides on `Projector` operations, preserving the constructor default and each call's geometry graph without rebuilding the projector (#29).
|
|
15
|
+
- Angle-tensor inputs for circular trajectory helpers and differentiable `transform_trajectory` rigid transforms (#29).
|
|
16
|
+
- `beam="parallel3d"` and `circular_trajectory_3d_parallel`, using infinite parallel rays and the PyTorch backend on CPU and CUDA, with data and first-order geometry gradients (#29).
|
|
17
|
+
- `Projector(native_volume_layout=True)` returns native full-CUDA cone backprojections of unbatched flat-detector inputs as `(D, H, W)` tensors with the kernels' `(W, H, D)`-contiguous strides instead of contiguous copies, and native cone projections accept such tensors without copying. An iterative solver that keeps its volumes in this layout saves two whole-volume copies per normal-operator application on one device. The default is unchanged (#27).
|
|
18
|
+
|
|
19
|
+
### Changed
|
|
20
|
+
|
|
21
|
+
- The native cone backprojection kernel steps through the z crossings of each xy cell in an inner loop and groups adjacent detector rows in a warp, so the atomic additions of a warp hit neighbouring voxels of one column. Each ray's segments and arithmetic are unchanged; backprojection takes 1.25–1.8x less time on large cone and helical problems, depending on the geometry (#35).
|
|
22
|
+
- Projector execution planning takes no allocator snapshot when a full launch clearly fits: lower bounds of the memory limits come from `mem_get_info` and one statistics query, and the exact planner runs only when they do not suffice. Plans are unchanged. Host time per call for small problems drops by about 1 ms. Distributed operators keep the exact planner (#26).
|
|
23
|
+
- Distributed calls and solvers check rank agreement with one fixed-size collective per request instead of Python object exchanges, skip the callback protocol when no callback is given, and accumulate solver norms on the device with one host read per norm. Results are unchanged (#34).
|
|
24
|
+
- The introduction video explains 2.3.0 tiled execution and view, space and hybrid partitions with a computed walnut example; the README GIF is 12 MB instead of 22 MB (#25).
|
|
25
|
+
|
|
26
|
+
### Fixed
|
|
27
|
+
|
|
28
|
+
- Detector surfaces expand world points on each compute device after view sharding, preserving native composition precision. Memory planning uses local view counts and composition dtypes. Streamed execution accepts CPU or CUDA offsets from both surface interfaces (#30).
|
|
29
|
+
- Surface snapshots preserve gradient connections on CPU. Streamed composition and geometry backward use bounded fragments. Distributed calls check callback layouts and propagate validation errors before numerical work (#30).
|
|
30
|
+
- Native surface admission sizes validation batches from available memory and staging limits, replacing fixed one-view batches (#32).
|
|
31
|
+
- Rigid trajectory transforms accept valid surface detector frames whose reference center coincides with the source or whose reference plane is edge-on (#29).
|
|
32
|
+
- Call-time trajectory validation checks flat detectors by corners and coordinate bounds, with exact pixel checks for ambiguous views. Frames stage once per validation device. Native surface validation uses bounded batches on each compute device after view sharding (#29).
|
|
33
|
+
|
|
8
34
|
## [2.3.0] - 2026-10-14
|
|
9
35
|
|
|
10
36
|
### Added
|
|
@@ -13,8 +13,8 @@ authors:
|
|
|
13
13
|
given-names: Yipeng
|
|
14
14
|
email: yipeng.sun@fau.de
|
|
15
15
|
alias: sypsyp97
|
|
16
|
-
version: 2.
|
|
17
|
-
date-released: "2026-10-
|
|
16
|
+
version: 2.4.0
|
|
17
|
+
date-released: "2026-10-20"
|
|
18
18
|
license: Apache-2.0
|
|
19
19
|
repository-code: "https://github.com/sypsyp97/diffct"
|
|
20
20
|
url: "https://sypsyp97.github.io/diffct/"
|
|
@@ -43,9 +43,9 @@ preferred-citation:
|
|
|
43
43
|
- family-names: Maier
|
|
44
44
|
given-names: Andreas
|
|
45
45
|
year: 2026
|
|
46
|
-
month:
|
|
46
|
+
month: 10
|
|
47
47
|
journal: Preprints
|
|
48
48
|
publisher:
|
|
49
49
|
name: Preprints
|
|
50
|
-
doi: 10.20944/preprints202605.1446.
|
|
51
|
-
url: "https://doi.org/10.20944/preprints202605.1446.
|
|
50
|
+
doi: 10.20944/preprints202605.1446.v2
|
|
51
|
+
url: "https://doi.org/10.20944/preprints202605.1446.v2"
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.5
|
|
2
2
|
Name: diffct
|
|
3
|
-
Version: 2.
|
|
3
|
+
Version: 2.4.0
|
|
4
4
|
Summary: Differentiable CT projection and reconstruction: precision modes, streaming, distributed execution and geometry gradients
|
|
5
5
|
Project-URL: Homepage, https://github.com/sypsyp97/diffct
|
|
6
6
|
Project-URL: Documentation, https://sypsyp97.github.io/diffct/
|
|
@@ -52,8 +52,8 @@ Description-Content-Type: text/markdown
|
|
|
52
52
|
|
|
53
53
|
<p align="center">
|
|
54
54
|
<a href="https://github.com/sypsyp97/diffct/blob/main/docs/assets/diffct_intro.mp4">Intro video (MP4)</a> ·
|
|
55
|
-
<a href="https://www.preprints.org/manuscript/202605.1446/
|
|
56
|
-
<a href="https://doi.org/10.20944/preprints202605.1446.
|
|
55
|
+
<a href="https://www.preprints.org/manuscript/202605.1446/v2">Technical report</a> ·
|
|
56
|
+
<a href="https://doi.org/10.20944/preprints202605.1446.v2">DOI</a>
|
|
57
57
|
</p>
|
|
58
58
|
|
|
59
59
|
<p align="center">
|
|
@@ -68,16 +68,18 @@ Description-Content-Type: text/markdown
|
|
|
68
68
|
> **Version 2.0** replaces the circular-orbit API of 1.x with per-view trajectories. Code written for 1.x needs changes; see the [migration notes](https://github.com/sypsyp97/diffct/blob/main/docs/MIGRATION.md).
|
|
69
69
|
> For Apple Silicon, use the [MLX port by Linda-Sophie Schneider](https://github.com/Linda-SophieSchneider/DiffCT-MLX).
|
|
70
70
|
|
|
71
|
+
This README describes the source checkout, including [unreleased APIs](CHANGELOG.md#unreleased). Per-call trajectories, angle-tensor circular helpers, `transform_trajectory`, `parallel3d` and `native_volume_layout` require installation from source.
|
|
72
|
+
|
|
71
73
|
## Why diffct
|
|
72
74
|
|
|
73
75
|
To our knowledge, diffct is the only open-source GPU CT library that combines arbitrary per-view trajectories, autograd through the volume, first-order gradients through the acquisition geometry, and built-in multi-GPU and multi-node execution. We compared it with LEAP, TIGRE, ASTRA/tomosipo, DiffDRR and others in October 2026.
|
|
74
76
|
|
|
75
|
-
- **Arbitrary trajectories.** Each view has its own source, detector centre and detector axes. Circular, helical, saddle, sinusoidal, random and calibrated scans use the same code. Trajectory tensors with `requires_grad=True` receive geometry gradients for calibration.
|
|
77
|
+
- **Arbitrary trajectories.** Each view has its own ray direction or source, detector centre and detector axes. Circular, helical, saddle, sinusoidal, random and calibrated scans use the same code. Trajectory tensors with `requires_grad=True` receive geometry gradients for calibration.
|
|
76
78
|
- **Curved detectors.** `ParameterizedSurface` defines native arc, cylindrical or arbitrary pixels from global coordinates and explicit parameters, sampled in bounded view/pixel batches. Legacy `detector_surface(u, v)` callbacks remain available. Projection and matched backprojection support first-order surface gradients. See [the surface API](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#parameterized-detector-surfaces).
|
|
77
79
|
- **Matched operators.** `project()` and `backproject()` form an exact adjoint pair for the cell-constant Siddon model. Both support PyTorch autograd, with volume/sinogram gradients and Hessian-vector products.
|
|
78
80
|
- **Large volumes.** Automatic spatial, view and pixel blocks with bounded transfer buffers control working memory. `project_into`/`backproject_into` accept caller-owned tensor or disk stores; tensor autograd keeps complete data on its input device. See [chunked execution](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#chunked-execution).
|
|
79
81
|
- **Multi-GPU and multi-node.** Use `devices=[0, 1, 2, 3]` in one process, or torchrun and NCCL across ranks. Views, space and hybrid partitions control replication of volumes and projections. Speedup depends on workload and communication costs.
|
|
80
|
-
- **Precision and training.** Explicit native/PyTorch backends, FP32/FP64 modes, per-axis voxel spacing, leading batch/channel dimensions and data `torch.func` transforms. The portable path supports CPU and PyTorch CUDA/ROCm devices; see [execution contracts
|
|
82
|
+
- **Precision and training.** Explicit native/PyTorch backends, FP32/FP64 modes, per-axis voxel spacing, leading batch/channel dimensions and data `torch.func` transforms. The portable path supports CPU and PyTorch CUDA/ROCm devices; see [execution contracts](docs/EXECUTION.md).
|
|
81
83
|
- **Solver API.** `cgls`, `sirt` and `tv_reconstruction` report residual history, tolerances and callbacks, with scan/data fingerprints and complete checkpoint state. `cgls_into` keeps state in bounded mapped blocks.
|
|
82
84
|
- **Differentiable FBP and FDK.** `fbp`, `fdk` and approximate `helical_fdk` compose the public filters, weights and analytical backprojectors. Their backward kernels transpose the actual interpolation stencil; custom filters and weights can be learned. See the [analytical guide](docs/source/analytical_learning.rst).
|
|
83
85
|
|
|
@@ -99,7 +101,7 @@ From source, with the examples:
|
|
|
99
101
|
```bash
|
|
100
102
|
git clone https://github.com/sypsyp97/diffct.git
|
|
101
103
|
cd diffct && pip install -e ".[cu12]"
|
|
102
|
-
python examples/quickstart.py #
|
|
104
|
+
python examples/quickstart.py # prints adjoint mismatch and gradients for each beam
|
|
103
105
|
```
|
|
104
106
|
|
|
105
107
|
<details>
|
|
@@ -138,7 +140,11 @@ loss.backward() # x.grad = A^T (A x - y)
|
|
|
138
140
|
|
|
139
141
|
**Custom or calibrated trajectories.** Pass calibrated tensors directly instead of a generator; no circular fit is needed. See [REFERENCE](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#custom-or-calibrated-trajectories).
|
|
140
142
|
|
|
141
|
-
**Geometry gradients.**
|
|
143
|
+
**Geometry gradients.** Reuse one projector with `A.project(volume, trajectory=current)` or `A.backproject(sinogram, trajectory=current)`. Circular helpers accept learnable angle tensors, and `transform_trajectory` applies differentiable rigid transforms. Omitting the override keeps the constructor trajectory. See [REFERENCE](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#geometry-gradients).
|
|
144
|
+
|
|
145
|
+
**3D parallel beams.** Use `circular_trajectory_3d_parallel(...)` and `beam="parallel3d"`. The PyTorch backend provides matched projection and backprojection on CPU and CUDA.
|
|
146
|
+
|
|
147
|
+
**Native volume layout.** For unbatched, flat-detector cone operations, set `native_volume_layout=True` to avoid volume layout copies. This applies to the native full-CUDA path on one compute device. Shapes stay `(D, H, W)`; the returned strides follow the kernel layout. The default is `False`. See [chunked execution](docs/REFERENCE.md#chunked-execution) for limits.
|
|
142
148
|
|
|
143
149
|
**Several GPUs and nodes.** See [REFERENCE](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#several-gpus-and-nodes) and [DISTRIBUTED.md](https://github.com/sypsyp97/diffct/blob/main/docs/DISTRIBUTED.md).
|
|
144
150
|
|
|
@@ -171,18 +177,18 @@ volume slabs use `--partition space` under torchrun. Run
|
|
|
171
177
|
|
|
172
178
|
<p align="center">
|
|
173
179
|
<picture>
|
|
174
|
-
<source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_dark.png?v=2.
|
|
175
|
-
<img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_light.png?v=2.
|
|
180
|
+
<source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_dark.png?v=2.4.0">
|
|
181
|
+
<img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_light.png?v=2.4.0" width="100%" alt="Time per forward, adjoint and CGLS iteration on 1, 4 and 8 GPUs, 512^3 volume">
|
|
176
182
|
</picture>
|
|
177
183
|
</p>
|
|
178
184
|
|
|
179
185
|
| 512³ volume | 1 GPU | 4 GPUs, 1 node | 8 GPUs, 2 nodes |
|
|
180
186
|
|---|---:|---:|---:|
|
|
181
|
-
| Forward projection | 493.
|
|
182
|
-
| Adjoint (backprojection) |
|
|
183
|
-
| CGLS iteration |
|
|
187
|
+
| Forward projection | 493.40 ms | 131.78 ms (3.7×) | 72.42 ms (6.8×) |
|
|
188
|
+
| Adjoint (backprojection) | 1118.98 ms | 293.01 ms (3.8×) | 232.50 ms (4.8×) |
|
|
189
|
+
| CGLS iteration | 1744.92 ms | 465.62 ms (3.7×) | 337.98 ms (5.2×) |
|
|
184
190
|
|
|
185
|
-
diffct 2.
|
|
191
|
+
diffct 2.4.0, FP32, A100 64 GB GPUs; circular cone beam, 360 views, 1024×768 detector, pitch 1.25. Default memory settings, without an imposed allocator cap.
|
|
186
192
|
|
|
187
193
|
Projection times are medians of seven samples after three warmups. CGLS uses three ten-iteration solves after a two-iteration warmup; the median includes initialization and is divided by ten. Timings use synchronized wall time and the slowest rank. Raw samples and configuration: [scaling.json](https://github.com/sypsyp97/diffct/blob/main/docs/assets/scaling.json).
|
|
188
194
|
|
|
@@ -222,7 +228,7 @@ Walnut data: Meaney 2022, Zenodo 6986012, CC BY 4.0; see [examples/data/NOTICE](
|
|
|
222
228
|
|
|
223
229
|
## Examples
|
|
224
230
|
|
|
225
|
-
`quickstart.py`, `curved_detector.py`, `chunked_reconstruction.py`, `disk_reconstruction.py`, `benchmark_execution.py`, `analytical_reconstruction.py`, `learnable_reconstruction.py`, `helical_fdk.py`, `iterative_reconstruction.py`, `walnut_reconstruction.py`, `geometry_calibration.py`, `benchmark_projector.py`, `benchmark_scaling.py`, `plot_trajectory.py`, and the Slurm template `slurm/multi_node.sbatch`. Launch modes and distributed-loss rules are in [examples/README.md](https://github.com/sypsyp97/diffct/blob/main/examples/README.md).
|
|
231
|
+
`quickstart.py`, `curved_detector.py`, `chunked_reconstruction.py`, `disk_reconstruction.py`, `benchmark_execution.py`, `benchmark_trajectory_admission.py`, `analytical_reconstruction.py`, `learnable_reconstruction.py`, `helical_fdk.py`, `iterative_reconstruction.py`, `walnut_reconstruction.py`, `geometry_calibration.py`, `benchmark_projector.py`, `benchmark_scaling.py`, `plot_trajectory.py`, and the Slurm template `slurm/multi_node.sbatch`. Launch modes and distributed-loss rules are in [examples/README.md](https://github.com/sypsyp97/diffct/blob/main/examples/README.md).
|
|
226
232
|
|
|
227
233
|
## Documentation
|
|
228
234
|
|
|
@@ -253,12 +259,12 @@ If you use this library in your research, please cite the software and the techn
|
|
|
253
259
|
}
|
|
254
260
|
|
|
255
261
|
@article{202605.1446,
|
|
256
|
-
doi = {10.20944/preprints202605.1446.
|
|
257
|
-
url = {https://doi.org/10.20944/preprints202605.1446.
|
|
262
|
+
doi = {10.20944/preprints202605.1446.v2},
|
|
263
|
+
url = {https://doi.org/10.20944/preprints202605.1446.v2},
|
|
258
264
|
year = 2026,
|
|
259
|
-
month = {
|
|
265
|
+
month = {Oct},
|
|
260
266
|
publisher = {Preprints},
|
|
261
|
-
author = {Yipeng Sun and Linda-Sophie Schneider and Chengze
|
|
267
|
+
author = {Yipeng Sun and Linda-Sophie Schneider and Chengze Ye and Andreas Maier},
|
|
262
268
|
title = {diffct: Differentiable CT Operators from Circular Orbits to Arbitrary Trajectories},
|
|
263
269
|
journal = {Preprints}
|
|
264
270
|
}
|
|
@@ -21,8 +21,8 @@
|
|
|
21
21
|
|
|
22
22
|
<p align="center">
|
|
23
23
|
<a href="https://github.com/sypsyp97/diffct/blob/main/docs/assets/diffct_intro.mp4">Intro video (MP4)</a> ·
|
|
24
|
-
<a href="https://www.preprints.org/manuscript/202605.1446/
|
|
25
|
-
<a href="https://doi.org/10.20944/preprints202605.1446.
|
|
24
|
+
<a href="https://www.preprints.org/manuscript/202605.1446/v2">Technical report</a> ·
|
|
25
|
+
<a href="https://doi.org/10.20944/preprints202605.1446.v2">DOI</a>
|
|
26
26
|
</p>
|
|
27
27
|
|
|
28
28
|
<p align="center">
|
|
@@ -37,16 +37,18 @@
|
|
|
37
37
|
> **Version 2.0** replaces the circular-orbit API of 1.x with per-view trajectories. Code written for 1.x needs changes; see the [migration notes](https://github.com/sypsyp97/diffct/blob/main/docs/MIGRATION.md).
|
|
38
38
|
> For Apple Silicon, use the [MLX port by Linda-Sophie Schneider](https://github.com/Linda-SophieSchneider/DiffCT-MLX).
|
|
39
39
|
|
|
40
|
+
This README describes the source checkout, including [unreleased APIs](CHANGELOG.md#unreleased). Per-call trajectories, angle-tensor circular helpers, `transform_trajectory`, `parallel3d` and `native_volume_layout` require installation from source.
|
|
41
|
+
|
|
40
42
|
## Why diffct
|
|
41
43
|
|
|
42
44
|
To our knowledge, diffct is the only open-source GPU CT library that combines arbitrary per-view trajectories, autograd through the volume, first-order gradients through the acquisition geometry, and built-in multi-GPU and multi-node execution. We compared it with LEAP, TIGRE, ASTRA/tomosipo, DiffDRR and others in October 2026.
|
|
43
45
|
|
|
44
|
-
- **Arbitrary trajectories.** Each view has its own source, detector centre and detector axes. Circular, helical, saddle, sinusoidal, random and calibrated scans use the same code. Trajectory tensors with `requires_grad=True` receive geometry gradients for calibration.
|
|
46
|
+
- **Arbitrary trajectories.** Each view has its own ray direction or source, detector centre and detector axes. Circular, helical, saddle, sinusoidal, random and calibrated scans use the same code. Trajectory tensors with `requires_grad=True` receive geometry gradients for calibration.
|
|
45
47
|
- **Curved detectors.** `ParameterizedSurface` defines native arc, cylindrical or arbitrary pixels from global coordinates and explicit parameters, sampled in bounded view/pixel batches. Legacy `detector_surface(u, v)` callbacks remain available. Projection and matched backprojection support first-order surface gradients. See [the surface API](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#parameterized-detector-surfaces).
|
|
46
48
|
- **Matched operators.** `project()` and `backproject()` form an exact adjoint pair for the cell-constant Siddon model. Both support PyTorch autograd, with volume/sinogram gradients and Hessian-vector products.
|
|
47
49
|
- **Large volumes.** Automatic spatial, view and pixel blocks with bounded transfer buffers control working memory. `project_into`/`backproject_into` accept caller-owned tensor or disk stores; tensor autograd keeps complete data on its input device. See [chunked execution](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#chunked-execution).
|
|
48
50
|
- **Multi-GPU and multi-node.** Use `devices=[0, 1, 2, 3]` in one process, or torchrun and NCCL across ranks. Views, space and hybrid partitions control replication of volumes and projections. Speedup depends on workload and communication costs.
|
|
49
|
-
- **Precision and training.** Explicit native/PyTorch backends, FP32/FP64 modes, per-axis voxel spacing, leading batch/channel dimensions and data `torch.func` transforms. The portable path supports CPU and PyTorch CUDA/ROCm devices; see [execution contracts
|
|
51
|
+
- **Precision and training.** Explicit native/PyTorch backends, FP32/FP64 modes, per-axis voxel spacing, leading batch/channel dimensions and data `torch.func` transforms. The portable path supports CPU and PyTorch CUDA/ROCm devices; see [execution contracts](docs/EXECUTION.md).
|
|
50
52
|
- **Solver API.** `cgls`, `sirt` and `tv_reconstruction` report residual history, tolerances and callbacks, with scan/data fingerprints and complete checkpoint state. `cgls_into` keeps state in bounded mapped blocks.
|
|
51
53
|
- **Differentiable FBP and FDK.** `fbp`, `fdk` and approximate `helical_fdk` compose the public filters, weights and analytical backprojectors. Their backward kernels transpose the actual interpolation stencil; custom filters and weights can be learned. See the [analytical guide](docs/source/analytical_learning.rst).
|
|
52
54
|
|
|
@@ -68,7 +70,7 @@ From source, with the examples:
|
|
|
68
70
|
```bash
|
|
69
71
|
git clone https://github.com/sypsyp97/diffct.git
|
|
70
72
|
cd diffct && pip install -e ".[cu12]"
|
|
71
|
-
python examples/quickstart.py #
|
|
73
|
+
python examples/quickstart.py # prints adjoint mismatch and gradients for each beam
|
|
72
74
|
```
|
|
73
75
|
|
|
74
76
|
<details>
|
|
@@ -107,7 +109,11 @@ loss.backward() # x.grad = A^T (A x - y)
|
|
|
107
109
|
|
|
108
110
|
**Custom or calibrated trajectories.** Pass calibrated tensors directly instead of a generator; no circular fit is needed. See [REFERENCE](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#custom-or-calibrated-trajectories).
|
|
109
111
|
|
|
110
|
-
**Geometry gradients.**
|
|
112
|
+
**Geometry gradients.** Reuse one projector with `A.project(volume, trajectory=current)` or `A.backproject(sinogram, trajectory=current)`. Circular helpers accept learnable angle tensors, and `transform_trajectory` applies differentiable rigid transforms. Omitting the override keeps the constructor trajectory. See [REFERENCE](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#geometry-gradients).
|
|
113
|
+
|
|
114
|
+
**3D parallel beams.** Use `circular_trajectory_3d_parallel(...)` and `beam="parallel3d"`. The PyTorch backend provides matched projection and backprojection on CPU and CUDA.
|
|
115
|
+
|
|
116
|
+
**Native volume layout.** For unbatched, flat-detector cone operations, set `native_volume_layout=True` to avoid volume layout copies. This applies to the native full-CUDA path on one compute device. Shapes stay `(D, H, W)`; the returned strides follow the kernel layout. The default is `False`. See [chunked execution](docs/REFERENCE.md#chunked-execution) for limits.
|
|
111
117
|
|
|
112
118
|
**Several GPUs and nodes.** See [REFERENCE](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#several-gpus-and-nodes) and [DISTRIBUTED.md](https://github.com/sypsyp97/diffct/blob/main/docs/DISTRIBUTED.md).
|
|
113
119
|
|
|
@@ -140,18 +146,18 @@ volume slabs use `--partition space` under torchrun. Run
|
|
|
140
146
|
|
|
141
147
|
<p align="center">
|
|
142
148
|
<picture>
|
|
143
|
-
<source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_dark.png?v=2.
|
|
144
|
-
<img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_light.png?v=2.
|
|
149
|
+
<source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_dark.png?v=2.4.0">
|
|
150
|
+
<img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_light.png?v=2.4.0" width="100%" alt="Time per forward, adjoint and CGLS iteration on 1, 4 and 8 GPUs, 512^3 volume">
|
|
145
151
|
</picture>
|
|
146
152
|
</p>
|
|
147
153
|
|
|
148
154
|
| 512³ volume | 1 GPU | 4 GPUs, 1 node | 8 GPUs, 2 nodes |
|
|
149
155
|
|---|---:|---:|---:|
|
|
150
|
-
| Forward projection | 493.
|
|
151
|
-
| Adjoint (backprojection) |
|
|
152
|
-
| CGLS iteration |
|
|
156
|
+
| Forward projection | 493.40 ms | 131.78 ms (3.7×) | 72.42 ms (6.8×) |
|
|
157
|
+
| Adjoint (backprojection) | 1118.98 ms | 293.01 ms (3.8×) | 232.50 ms (4.8×) |
|
|
158
|
+
| CGLS iteration | 1744.92 ms | 465.62 ms (3.7×) | 337.98 ms (5.2×) |
|
|
153
159
|
|
|
154
|
-
diffct 2.
|
|
160
|
+
diffct 2.4.0, FP32, A100 64 GB GPUs; circular cone beam, 360 views, 1024×768 detector, pitch 1.25. Default memory settings, without an imposed allocator cap.
|
|
155
161
|
|
|
156
162
|
Projection times are medians of seven samples after three warmups. CGLS uses three ten-iteration solves after a two-iteration warmup; the median includes initialization and is divided by ten. Timings use synchronized wall time and the slowest rank. Raw samples and configuration: [scaling.json](https://github.com/sypsyp97/diffct/blob/main/docs/assets/scaling.json).
|
|
157
163
|
|
|
@@ -191,7 +197,7 @@ Walnut data: Meaney 2022, Zenodo 6986012, CC BY 4.0; see [examples/data/NOTICE](
|
|
|
191
197
|
|
|
192
198
|
## Examples
|
|
193
199
|
|
|
194
|
-
`quickstart.py`, `curved_detector.py`, `chunked_reconstruction.py`, `disk_reconstruction.py`, `benchmark_execution.py`, `analytical_reconstruction.py`, `learnable_reconstruction.py`, `helical_fdk.py`, `iterative_reconstruction.py`, `walnut_reconstruction.py`, `geometry_calibration.py`, `benchmark_projector.py`, `benchmark_scaling.py`, `plot_trajectory.py`, and the Slurm template `slurm/multi_node.sbatch`. Launch modes and distributed-loss rules are in [examples/README.md](https://github.com/sypsyp97/diffct/blob/main/examples/README.md).
|
|
200
|
+
`quickstart.py`, `curved_detector.py`, `chunked_reconstruction.py`, `disk_reconstruction.py`, `benchmark_execution.py`, `benchmark_trajectory_admission.py`, `analytical_reconstruction.py`, `learnable_reconstruction.py`, `helical_fdk.py`, `iterative_reconstruction.py`, `walnut_reconstruction.py`, `geometry_calibration.py`, `benchmark_projector.py`, `benchmark_scaling.py`, `plot_trajectory.py`, and the Slurm template `slurm/multi_node.sbatch`. Launch modes and distributed-loss rules are in [examples/README.md](https://github.com/sypsyp97/diffct/blob/main/examples/README.md).
|
|
195
201
|
|
|
196
202
|
## Documentation
|
|
197
203
|
|
|
@@ -222,12 +228,12 @@ If you use this library in your research, please cite the software and the techn
|
|
|
222
228
|
}
|
|
223
229
|
|
|
224
230
|
@article{202605.1446,
|
|
225
|
-
doi = {10.20944/preprints202605.1446.
|
|
226
|
-
url = {https://doi.org/10.20944/preprints202605.1446.
|
|
231
|
+
doi = {10.20944/preprints202605.1446.v2},
|
|
232
|
+
url = {https://doi.org/10.20944/preprints202605.1446.v2},
|
|
227
233
|
year = 2026,
|
|
228
|
-
month = {
|
|
234
|
+
month = {Oct},
|
|
229
235
|
publisher = {Preprints},
|
|
230
|
-
author = {Yipeng Sun and Linda-Sophie Schneider and Chengze
|
|
236
|
+
author = {Yipeng Sun and Linda-Sophie Schneider and Chengze Ye and Andreas Maier},
|
|
231
237
|
title = {diffct: Differentiable CT Operators from Circular Orbits to Arbitrary Trajectories},
|
|
232
238
|
journal = {Preprints}
|
|
233
239
|
}
|
|
@@ -30,6 +30,8 @@ from .reconstruction import (
|
|
|
30
30
|
from .solvers import SolverResult, cgls, sirt, tv_reconstruction, cgls_into, total_variation
|
|
31
31
|
|
|
32
32
|
from .geometry import (
|
|
33
|
+
transform_trajectory,
|
|
34
|
+
circular_trajectory_3d_parallel,
|
|
33
35
|
circular_trajectory_3d,
|
|
34
36
|
random_trajectory_3d,
|
|
35
37
|
spiral_trajectory_3d,
|
|
@@ -62,7 +64,7 @@ try:
|
|
|
62
64
|
except ImportError:
|
|
63
65
|
pass
|
|
64
66
|
|
|
65
|
-
__version__ = '2.
|
|
67
|
+
__version__ = '2.4.0'
|
|
66
68
|
|
|
67
69
|
__all__ = [
|
|
68
70
|
'Projector',
|
|
@@ -88,6 +90,8 @@ __all__ = [
|
|
|
88
90
|
'FanBackprojectorFunction',
|
|
89
91
|
'ConeProjectorFunction',
|
|
90
92
|
'ConeBackprojectorFunction',
|
|
93
|
+
'transform_trajectory',
|
|
94
|
+
'circular_trajectory_3d_parallel',
|
|
91
95
|
'circular_trajectory_3d',
|
|
92
96
|
'random_trajectory_3d',
|
|
93
97
|
'spiral_trajectory_3d',
|
|
@@ -11,6 +11,7 @@ import torch.distributed as dist
|
|
|
11
11
|
|
|
12
12
|
from .pipeline import _fragments
|
|
13
13
|
from .storage import NpyStore, TensorStore
|
|
14
|
+
from .distribution import control_status, preflight
|
|
14
15
|
|
|
15
16
|
|
|
16
17
|
def count(value, name, positive=False):
|
|
@@ -25,47 +26,56 @@ def finite(value, name, positive=False):
|
|
|
25
26
|
return float(value)
|
|
26
27
|
|
|
27
28
|
|
|
28
|
-
def agreement(operator, error=None, token=None):
|
|
29
|
+
def agreement(operator, error=None, token=None, *, cache_key=None):
|
|
29
30
|
"""Every rank rejects a bad request before numerical execution starts."""
|
|
30
31
|
world = getattr(operator, "world_size", 1)
|
|
31
32
|
if world == 1:
|
|
32
33
|
if error is not None:
|
|
33
34
|
raise error
|
|
34
35
|
return
|
|
35
|
-
|
|
36
|
-
|
|
37
|
-
|
|
38
|
-
|
|
39
|
-
|
|
40
|
-
|
|
36
|
+
error = None if error is None else ValueError("distributed solver agreement failed: " + str(error))
|
|
37
|
+
mismatch = "distributed solver agreement failed: inconsistent checkpoint iteration or controls"
|
|
38
|
+
partition = getattr(operator, "_partition", None)
|
|
39
|
+
if partition is not None:
|
|
40
|
+
partition.preflight(error, token, cache_key=cache_key, mismatch=mismatch)
|
|
41
|
+
else:
|
|
42
|
+
preflight(error, token, process_group=operator.process_group, world_size=world,
|
|
43
|
+
mismatch=mismatch)
|
|
41
44
|
|
|
42
45
|
|
|
43
|
-
def
|
|
44
|
-
"""
|
|
45
|
-
world = getattr(operator, "world_size", 1)
|
|
46
|
+
def owns_values(operator, kind):
|
|
47
|
+
"""Select one copy of each physical ray or voxel for scalar reductions."""
|
|
46
48
|
partition = getattr(operator, "partition", "views")
|
|
47
49
|
rank = getattr(operator, "rank", 0)
|
|
48
|
-
|
|
50
|
+
return (getattr(operator, "view_rank", 0) == 0 if kind == "volume"
|
|
49
51
|
else getattr(operator, "space_rank", 0) == 0) if partition == "hybrid" else (
|
|
50
52
|
rank == 0 if (kind == "volume") == (partition == "views") else True
|
|
51
53
|
)
|
|
52
|
-
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
def scalar_sum(operator, value, kind):
|
|
57
|
+
"""Count physical rays/voxels once, including view/space replicas."""
|
|
58
|
+
world = getattr(operator, "world_size", 1)
|
|
59
|
+
value = value if owns_values(operator, kind) else 0.
|
|
53
60
|
if world > 1:
|
|
54
61
|
group = operator.process_group
|
|
55
62
|
device = torch.device("cuda", torch.cuda.current_device()) if dist.get_backend(group) == "nccl" else torch.device("cpu")
|
|
56
|
-
accumulator = torch.
|
|
63
|
+
accumulator = torch.as_tensor(value, dtype=torch.float64, device=device).clone()
|
|
57
64
|
dist.all_reduce(accumulator, group=group)
|
|
58
|
-
|
|
59
|
-
return value
|
|
65
|
+
return accumulator.item()
|
|
66
|
+
return float(value)
|
|
60
67
|
|
|
61
68
|
|
|
62
69
|
def squared_norm(operator, value, kind, elements=1048576):
|
|
63
70
|
store = TensorStore(value) if isinstance(value, torch.Tensor) else value
|
|
64
71
|
total = 0.
|
|
65
|
-
|
|
66
|
-
|
|
67
|
-
|
|
68
|
-
|
|
72
|
+
if owns_values(operator, kind):
|
|
73
|
+
total = torch.zeros((), dtype=torch.float64, device=getattr(store, "device", "cpu"))
|
|
74
|
+
for selected in _fragments(tuple(store.shape), elements):
|
|
75
|
+
# Clone even an FP64 view: norms must never alter caller measurements.
|
|
76
|
+
block = store.read(selected).to(dtype=torch.float64).clone()
|
|
77
|
+
# Preserve the block reductions and their sequential FP64 order.
|
|
78
|
+
total.add_(block.square_().sum().to(device=total.device))
|
|
69
79
|
return scalar_sum(operator, total, kind)
|
|
70
80
|
|
|
71
81
|
|
|
@@ -270,15 +280,22 @@ def save_checkpoint(operator, directory, states, metadata, iteration, history, m
|
|
|
270
280
|
agreement(operator, error, iteration)
|
|
271
281
|
|
|
272
282
|
|
|
283
|
+
def callbacks_enabled(operator, callback):
|
|
284
|
+
"""Check callback presence once, including callbacks on only some ranks."""
|
|
285
|
+
return bool(control_status((int(callback is not None),),
|
|
286
|
+
process_group=getattr(operator, "process_group", None),
|
|
287
|
+
world_size=getattr(operator, "world_size", 1))[0])
|
|
288
|
+
|
|
289
|
+
|
|
273
290
|
def callback_stop(operator, callback, iteration, residual, volume):
|
|
274
291
|
error, stop = None, False
|
|
275
292
|
try:
|
|
276
293
|
stop = callback is not None and callback(iteration, residual, volume) is False
|
|
277
294
|
except Exception as caught:
|
|
278
295
|
error = caught
|
|
279
|
-
|
|
280
|
-
|
|
281
|
-
|
|
282
|
-
|
|
283
|
-
|
|
284
|
-
return stop
|
|
296
|
+
failed, stop = control_status((int(error is not None), int(stop)),
|
|
297
|
+
process_group=getattr(operator, "process_group", None),
|
|
298
|
+
world_size=getattr(operator, "world_size", 1))
|
|
299
|
+
if failed:
|
|
300
|
+
agreement(operator, error)
|
|
301
|
+
return bool(stop)
|