Make SynDiff run on current PyTorch, and fix the validation/checkpoint/inference paths

[mertcanozdemir]

SynDiff does not currently run on a modern PyTorch install, and a few defects
in the validation, checkpoint and inference paths survive into released code
because nothing exercises them. This branch fixes those, without touching the
training objective or the diffusion process.

Everything here was reproduced and then re-verified by actually running
train.py and test.py on the bundled sample data, on CPU and on a CUDA GPU
(PyTorch 2.13, NumPy 2.5, Python 3.12).

Nothing changes for the published commands. Every default was chosen to
reproduce the previous hard-coded value, and the multi-process path
(--num_process_per_node > 1) takes exactly the code path it took before.

---

The repository does not import

utils/op/fused_act.py and utils/op/upfirdn2d.py JIT-compiled their CUDA
extensions at import time. backbones imports these transitively, so on any
machine without ninja and a matching CUDA toolchain — including every CPU-only
install — nothing could be imported at all:

RuntimeError: Ninja is required to load C++ extensions

Both modules already contained pure-PyTorch fallbacks (upfirdn2d_native, and
the CPU branch of fused_leaky_relu); the unconditional load() simply ran
before they could ever be reached. They are now built inside a try/except
that warns once and routes to the native path.

Where the kernels do build, behaviour is unchanged — tests/test_ops.py
asserts the CUDA and native paths agree. On this machine the forward results
are bit-identical and gradients match to 1.2e-07.

Correctness fixes

The 2→1 validation direction was scored with the wrong network.
train_syndiff() runs two validation loops. The second swaps the batch to
(y_val, x_val) so the target becomes contrast 2 — but still sampled from
gen_diffusive_1, the network that produces contrast 1. Every value written to
val_l1_loss[1] and val_psnr_values[1] was therefore meaningless, and
gen_diffusive_2 was never validated. This does not affect trained weights
(validation is observational), but it does affect anything selected by reading
those numbers.

Training crashed at the very end of every run. The validation buffers were
allocated with args.num_epoch epoch slots while the loop runs over
range(init_epoch, args.num_epoch + 1):

File "train.py", line 711, in train_syndiff
  val_l1_loss[0,epoch,iteration]=abs(fake_sample1 -real_data).mean()
IndexError: index 1 is out of bounds for axis 1 with size 1

The final checkpoint is written before this point, so the loss is the last
epoch's metrics plus a traceback on every run.

test.py could not write its output. syn_im1/syn_im2 were preallocated
as (256, 256, N) while every sample is passed through
CenterCrop((256, 152)) two lines earlier. The crop size is fixed, so this
fired on the first test slice for every dataset:

ValueError: could not broadcast input array from shape (256,152) into shape (256,256)

Slices are now collected and stacked, so the stored volume matches the cropped
geometry, and the crop size is exposed as --crop_h/--crop_w (defaulting to
the previous values).

--resume is dead on PyTorch ≥ 2.6. content.pth stores the run's
argparse.Namespace next to the tensors, which the weights_only=True default
rejects:

_pickle.UnpicklingError: Weights only load failed.
... Unsupported global: GLOBAL argparse.Namespace

Checkpoint keys were truncated blindly. load_checkpoint() rewrote every
key as key[7:] on the assumption that a DDP wrapper had added module.. For
a checkpoint saved without one, that removes seven characters of a real
parameter name. Both loaders now strip the prefix only where present, and
repeatedly, so all three layouts load.

Several batch sizes crash the discriminators. The minibatch-stddev block
reshapes with group = min(batch, 4), which is only a valid view when group
divides the batch:

batch=5  RuntimeError: shape '[4, -1, 1, 64, 2, 2]' is invalid for input of size 1280
batch=6  RuntimeError: ...
batch=7  RuntimeError: ...

The group size now steps down to the largest divisor of the batch. Sizes that
worked before keep their previous grouping.

Padding silently missed the target size. LoadDataSet applied a single
symmetric int((256 - size) / 2) to both sides, so an odd difference lost a
pixel:

>>> LoadDataSet(<151-wide volume>).shape
(2, 1, 256, 255)

Nothing checked the result, so this only surfaced much later as a shape error
inside the network. Oversized inputs, a wrong variable name, a missing file and
mismatched slice counts between the two contrasts now each raise a message that
names the problem.

args.num_channels = 1 was assigned after both NCSNpp generators were
built and is not read by define_G, so it changed no network — but it was
serialised into content.pth, recording num_channels = 1 for a run whose
generators were built with 2.

Arguments that did not reach the code they name

--ngf was applied to the diffusive discriminators only; the translation
networks and cycle discriminators used define_G/define_D's own default of
64, so a run with a different --ngf got a model whose halves disagreed.
--image_size was passed to NCSNpp while CreateDatasetSynthesis padded to
a hard-coded 256. Both are threaded through. The published commands pass
--ngf 64 and --image_size 256, which are the previous constants.

Single-process and CPU runs

init_processes() called torch.cuda.set_device() and
dist.init_process_group(backend='nccl') unconditionally, so a single-GPU run
still had to stand up a NCCL group and a CPU-only run was impossible — which is
also why the training loop could not be tested without a GPU. A process group
is now created only when the run spans processes, and DDP wraps only when one
is live. size > 1 is untouched.

Wasted work in the discriminator phases

Both discriminator updates ran the generators with autograd enabled and never
detached the results, so errD_fake.backward() propagated through all four
generators and filled .grad buffers that the following zero_grad() threw
away. Wrapping those forwards in torch.no_grad() leaves the update
mathematically identical. Measured on one branch (256×256, batch 1, GPU, 10
iterations after warm-up):

	time	transient memory
no_grad	67.3 ms/iter	0.14 GB
before	156.4 ms/iter	0.56 GB

Peak memory for the whole step is unchanged — that peak is set by the generator
update, which legitimately keeps its graph.

Tests

The repository had no tests, which is why several of the defects above could
sit in released code: nothing ever checked the values they produced. The suite
runs on CPU in about five seconds.

pip install -r requirements.txt
python -m pytest tests/

• tests/test_diffusion.py — the variance-preserving identity, monotone signal
  decay, non-negative posterior variance, collapse to x_0 at t == 0, and
  that sample_posterior is deterministic there and stochastic elsewhere.
• tests/test_models.py — shapes and gradients for NCSN++, both discriminators
  (batch sizes 1–8) and the translation networks; NCSN++ must reduce a
  reconstruction loss over a few optimiser steps.
• tests/test_dataset.py — padding reaches the target size for odd and even
  widths, stays centred, normalises to [-1, 1], and each rejected input
  raises a message naming the problem.
• tests/test_checkpoint.py — state_dicts load through all three module.
  layouts.
• tests/test_ops.py — the CUDA kernels agree with their native fallbacks,
  forward and backward; skipped where the extensions cannot build.

68 pass with the extensions built; 63 pass and 5 skip on a CPU-only install.

Housekeeping

diffusion.py collects the ~120 lines of sampler code that train.py and
test.py each carried a copy of. The two copies were verified byte-identical
before merging, and the definitions were moved verbatim — a change to the
diffusion process previously had to be made twice, or the two entry points
would silently disagree.

requirements.txt records the runtime set, which the README only described in
prose. torch>=1.13 is the floor because that release introduced
weights_only. ninja is listed separately since the fused kernels now fall
back without it.

The README's dependency list omitted numpy, h5py and scikit-image, all
imported at module scope. It also did not mention that each .mat file must
hold a variable named data_fs, or that the sample data under
SynDiff_sample_data/ is two raw 25-slice volumes (T1.mat, T2.mat) rather
than the data_<phase>_<contrast>.mat layout the loader expects — so it does
not run as shipped.

backbones/im2im.py and utils/utils.py are removed. Neither is referenced
anywhere. im2im.py holds a second copy of define_G/init_net, so the
DataParallel fix in this branch would not have applied to it; utils/utils.py
imports TensorFlow at module scope for two functions nothing calls. Happy to
drop this commit if you would rather keep them.

Not addressed

• The diffusion mathematics was reviewed by reading and is unchanged here; it
  was not re-derived against the paper, and no model was trained to
  convergence or compared against the released pretrained weights.
• --num_epoch 0 still raises ZeroDivisionError from
  CosineAnnealingLR(T_max=0). Degenerate input, left alone.

[mertcanozdemir]

claude fable detaylı analiz sonucu. commitleri ayrı ayrı oluşturmasını istedim. dilersenin kullanabilirsiniz.

-ENG-
Claude Fable's detailed analysis result. I asked it to create the commits separately. You can use it if you'd like.