From 75c1c761f0c6854023ed88ff678674fabd1829a5 Mon Sep 17 00:00:00 2001 From: Woody-Wan Date: Mon, 8 Jun 2026 16:46:32 -0500 Subject: [PATCH 1/3] Dmplug --- algo/dmplug.py | 110 ++++++++++++++++++++++++++++++++++ configs/algorithm/dmplug.yaml | 28 +++++++++ 2 files changed, 138 insertions(+) create mode 100644 algo/dmplug.py create mode 100644 configs/algorithm/dmplug.yaml diff --git a/algo/dmplug.py b/algo/dmplug.py new file mode 100644 index 0000000..a3c431e --- /dev/null +++ b/algo/dmplug.py @@ -0,0 +1,110 @@ +import copy +import json +import os + +import torch +from tqdm import tqdm +import torch.nn.functional as F +from torchvision.utils import save_image +from .base import Algo +from utils.scheduler import Scheduler +from utils.diffusion import DiffusionSampler + + +class DMPlug(Algo): + + ''' + DMPlug algorithm implemented in EDM framework. + ''' + + def __init__(self, + net, + forward_op, + diffusion_scheduler_config, + guidance_scale=1.0, + sde=False, + iteration=5000, + lr=0.1, + weight_decay=0.0, + loss_scaling='residual', + solver='euler', + grad_check_interval=0, + trace_interval=0): + super(DMPlug, self).__init__(net, forward_op) + self.net.eval().requires_grad_(False) + self.diffusion_scheduler_config = diffusion_scheduler_config + self.scheduler = Scheduler(**diffusion_scheduler_config) + self.guidance_scale = guidance_scale + self.sde = sde + self.iteration = iteration + self.lr = lr + self.weight_decay = weight_decay + self.loss_scaling = loss_scaling + self.solver = solver + self.grad_check_interval = grad_check_interval + self.trace_interval = trace_interval + if self.loss_scaling not in {'residual', 'mse', 'none'}: + raise ValueError("loss_scaling must be one of {'residual', 'mse', 'none'}") + + + def _measurement_numel(self, observation): + def numel(data): + if torch.is_tensor(data): + return data.numel() + size = getattr(data, "size", None) + return size() if callable(size) else size + + if torch.is_tensor(observation): + return observation.numel() + if isinstance(observation, (list, tuple)): + total = 0 + for item in observation: + data = getattr(item, "data", item) + total += numel(data) + return total + data = getattr(observation, "data", None) + if data is not None: + return numel(data) + raise TypeError("Cannot infer measurement size for loss_scaling='mse'.") + + def inference(self, observation, num_samples=1, **kwargs): + target = kwargs.get('target') + evaluator = kwargs.get('evaluator') + trace_dir = kwargs.get('trace_dir') + device = self.forward_op.device + if num_samples > 1: + if not torch.is_tensor(observation): + raise ValueError("DMPlug num_samples > 1 requires tensor observations.") + observation = observation.repeat(num_samples, *([1] * (observation.ndim - 1))) + x_initial = torch.randn(num_samples, self.net.img_channels, self.net.img_resolution, self.net.img_resolution, device=device) * self.scheduler.sigma_max + x_initial.requires_grad = True + + sampler = DiffusionSampler(self.scheduler, solver=self.solver) + pbar = tqdm(range(self.iteration)) + + optimizer = torch.optim.AdamW([x_initial], lr=self.lr, weight_decay=self.weight_decay) + + for iteration in pbar: + optimizer.zero_grad() + denoised = sampler.sample(self.net, x_initial, SDE=self.sde, verbose=False) + + gradient, loss_scale = self.forward_op.gradient(denoised, observation, return_loss=True) + + x_initial_grad = torch.autograd.grad( + outputs=denoised, + inputs=x_initial, + grad_outputs=gradient, + )[0] + if self.loss_scaling == 'residual': + x_initial_grad = x_initial_grad * 0.5 / torch.sqrt(loss_scale).clamp_min(1e-8) + elif self.loss_scaling == 'mse': + x_initial_grad = x_initial_grad / self._measurement_numel(observation) + x_initial.grad = x_initial_grad * self.guidance_scale + desc = f'Iteration {iteration + 1}/{self.iteration}. Data fitting loss: {display_loss}, x_initial.grad norm: {grad_norm}' + + optimizer.step() + pbar.set_description(desc) + + with torch.no_grad(): + denoised = sampler.sample(self.net, x_initial, SDE=self.sde, verbose=False) + return denoised diff --git a/configs/algorithm/dmplug.yaml b/configs/algorithm/dmplug.yaml new file mode 100644 index 0000000..bf3723c --- /dev/null +++ b/configs/algorithm/dmplug.yaml @@ -0,0 +1,28 @@ +name: DMPlug +method: + _target_: algo.dmplug.DMPlug + + diffusion_scheduler_config: + # DMPlug backpropagates through the full reverse sampler each optimizer step. + # Keep this much smaller than DPS's 1000 steps for cost, but 3 VP steps is + # numerically too coarse and leaves strong noise artifacts. + num_steps: 18 + schedule: 'vp' + timestep: 'vp' + scaling: 'vp' + + guidance_scale: 1.0 + lr: 1e-2 + weight_decay: 0.0 + iteration: 1000 + solver: euler + # mse: upstream DMPlug-style mean squared measurement error + # residual: DPS-style normalized sqrt(loss) update + # none: use forward_op.gradient exactly, matching inversebench's operator loss + # For multi-coil MRI, mse usually needs a much larger guidance_scale. + loss_scaling: mse + # Set to e.g. 100 to compare the manual VJP with direct objective.backward(). + grad_check_interval: 0 + # Save loss/eval histories and reconstructed tensors every N optimization iterations. + # Set to 0 to disable tracing. + trace_interval: 100 From 5c7f5685dcb10812bfe6c3493b22ad4f9c5bcce7 Mon Sep 17 00:00:00 2001 From: Woody-Wan Date: Mon, 8 Jun 2026 16:56:34 -0500 Subject: [PATCH 2/3] fix typos --- algo/dmplug.py | 5 +---- configs/algorithm/dmplug.yaml | 14 ++------------ 2 files changed, 3 insertions(+), 16 deletions(-) diff --git a/algo/dmplug.py b/algo/dmplug.py index a3c431e..2121953 100644 --- a/algo/dmplug.py +++ b/algo/dmplug.py @@ -28,8 +28,7 @@ def __init__(self, weight_decay=0.0, loss_scaling='residual', solver='euler', - grad_check_interval=0, - trace_interval=0): + ): super(DMPlug, self).__init__(net, forward_op) self.net.eval().requires_grad_(False) self.diffusion_scheduler_config = diffusion_scheduler_config @@ -41,8 +40,6 @@ def __init__(self, self.weight_decay = weight_decay self.loss_scaling = loss_scaling self.solver = solver - self.grad_check_interval = grad_check_interval - self.trace_interval = trace_interval if self.loss_scaling not in {'residual', 'mse', 'none'}: raise ValueError("loss_scaling must be one of {'residual', 'mse', 'none'}") diff --git a/configs/algorithm/dmplug.yaml b/configs/algorithm/dmplug.yaml index bf3723c..47db8dd 100644 --- a/configs/algorithm/dmplug.yaml +++ b/configs/algorithm/dmplug.yaml @@ -3,9 +3,6 @@ method: _target_: algo.dmplug.DMPlug diffusion_scheduler_config: - # DMPlug backpropagates through the full reverse sampler each optimizer step. - # Keep this much smaller than DPS's 1000 steps for cost, but 3 VP steps is - # numerically too coarse and leaves strong noise artifacts. num_steps: 18 schedule: 'vp' timestep: 'vp' @@ -16,13 +13,6 @@ method: weight_decay: 0.0 iteration: 1000 solver: euler - # mse: upstream DMPlug-style mean squared measurement error - # residual: DPS-style normalized sqrt(loss) update - # none: use forward_op.gradient exactly, matching inversebench's operator loss - # For multi-coil MRI, mse usually needs a much larger guidance_scale. + loss_scaling: mse - # Set to e.g. 100 to compare the manual VJP with direct objective.backward(). - grad_check_interval: 0 - # Save loss/eval histories and reconstructed tensors every N optimization iterations. - # Set to 0 to disable tracing. - trace_interval: 100 + From 8e3008bb4906bd54aa6b34a0bea59916e4a3b3fc Mon Sep 17 00:00:00 2001 From: devzhk Date: Tue, 7 Jul 2026 03:36:35 +0000 Subject: [PATCH 3/3] Fix DMPlug inference logging --- algo/dmplug.py | 36 ++++++++++++++++------------------- configs/algorithm/dmplug.yaml | 1 - 2 files changed, 16 insertions(+), 21 deletions(-) diff --git a/algo/dmplug.py b/algo/dmplug.py index 2121953..7ee7a71 100644 --- a/algo/dmplug.py +++ b/algo/dmplug.py @@ -1,23 +1,16 @@ -import copy -import json -import os - import torch from tqdm import tqdm -import torch.nn.functional as F -from torchvision.utils import save_image from .base import Algo from utils.scheduler import Scheduler from utils.diffusion import DiffusionSampler class DMPlug(Algo): - ''' DMPlug algorithm implemented in EDM framework. ''' - - def __init__(self, + + def __init__(self, net, forward_op, diffusion_scheduler_config, @@ -30,7 +23,7 @@ def __init__(self, solver='euler', ): super(DMPlug, self).__init__(net, forward_op) - self.net.eval().requires_grad_(False) + self.net.eval() self.diffusion_scheduler_config = diffusion_scheduler_config self.scheduler = Scheduler(**diffusion_scheduler_config) self.guidance_scale = guidance_scale @@ -63,11 +56,8 @@ def numel(data): if data is not None: return numel(data) raise TypeError("Cannot infer measurement size for loss_scaling='mse'.") - + def inference(self, observation, num_samples=1, **kwargs): - target = kwargs.get('target') - evaluator = kwargs.get('evaluator') - trace_dir = kwargs.get('trace_dir') device = self.forward_op.device if num_samples > 1: if not torch.is_tensor(observation): @@ -75,12 +65,12 @@ def inference(self, observation, num_samples=1, **kwargs): observation = observation.repeat(num_samples, *([1] * (observation.ndim - 1))) x_initial = torch.randn(num_samples, self.net.img_channels, self.net.img_resolution, self.net.img_resolution, device=device) * self.scheduler.sigma_max x_initial.requires_grad = True - + sampler = DiffusionSampler(self.scheduler, solver=self.solver) pbar = tqdm(range(self.iteration)) - + optimizer = torch.optim.AdamW([x_initial], lr=self.lr, weight_decay=self.weight_decay) - + for iteration in pbar: optimizer.zero_grad() denoised = sampler.sample(self.net, x_initial, SDE=self.sde, verbose=False) @@ -96,9 +86,15 @@ def inference(self, observation, num_samples=1, **kwargs): x_initial_grad = x_initial_grad * 0.5 / torch.sqrt(loss_scale).clamp_min(1e-8) elif self.loss_scaling == 'mse': x_initial_grad = x_initial_grad / self._measurement_numel(observation) - x_initial.grad = x_initial_grad * self.guidance_scale - desc = f'Iteration {iteration + 1}/{self.iteration}. Data fitting loss: {display_loss}, x_initial.grad norm: {grad_norm}' - + x_initial.grad = (x_initial_grad * self.guidance_scale).detach() + display_loss = torch.sqrt(loss_scale.detach()).item() + grad_norm = x_initial.grad.detach().norm().item() + desc = ( + f'Iteration {iteration + 1}/{self.iteration}. ' + f'Data fitting loss: {display_loss:.4f}, ' + f'x_initial.grad norm: {grad_norm:.4f}' + ) + optimizer.step() pbar.set_description(desc) diff --git a/configs/algorithm/dmplug.yaml b/configs/algorithm/dmplug.yaml index 47db8dd..25671ac 100644 --- a/configs/algorithm/dmplug.yaml +++ b/configs/algorithm/dmplug.yaml @@ -15,4 +15,3 @@ method: solver: euler loss_scaling: mse -