``.
+ update_rate : int
+ Write image / inline display every *update_rate* steps.
+ Scalar summaries are written every step.
+ """
- def close_plot(self):
- plt.close()
-
-class RECONmonitor:
- def __init__(self, recon_target, img_input):
- self.fig, self.axs = plt.subplots(2, 2, figsize=(16, 8))
+ def __init__(self, recon_target, img_input, logdir=None, update_rate=100):
self.recon_target = recon_target
- self.update_rate = 100
- self.img_input = img_input
+ self.img_input = img_input.astype(np.float32)
self.img_h, self.img_w = img_input.shape
- self.epoch = 0
- self.plot_y1, self.plot_y2 = [], []
- if self.recon_target in ["tomo", "tensor"]:
- self.plot_txt = "Sinogram"
- elif self.recon_target == "phase":
- self.plot_txt = "Intensity"
- self.__initial_plot()
-
- def __initial_plot(self):
- self.im0 = self.axs[0, 0].imshow(self.img_input, cmap="gray")
- self.axs[0, 0].set_title(self.plot_txt)
- self.fig.colorbar(self.im0, ax=self.axs[0, 0])
- self.axs[0, 0].set_aspect("equal", "box")
- self.im1 = self.axs[1, 0].imshow(self.img_input, cmap="jet")
- self.tx1 = self.axs[1, 0].set_title("SSIM map of " + self.plot_txt + " for iteration 0")
- self.fig.colorbar(self.im1, ax=self.axs[1, 0])
- self.axs[1, 0].set_aspect("equal")
- self.im2 = self.axs[0, 1].imshow(np.zeros((self.img_w, self.img_w)), cmap="gray")
- self.fig.colorbar(self.im2, ax=self.axs[0, 1])
- self.axs[0, 1].set_title("Reconstruction")
- (self.im3,) = self.axs[1, 1].plot([], [], "r-")
- self.axs[1, 1].set_title("Reconstruction loss")
- self.axs[1, 1].set_yscale("log")
- plt.tight_layout()
-
- def update_plot(self, step_result):
- self.epoch = self.epoch+1
- self.plot_x = np.arange(self.epoch)
- self.plot_y1.append(step_result['g_loss'].numpy())
- self.plot_y2.append(step_result['d_loss'].numpy())
-
- if (self.epoch + 1) % self.update_rate == 0:
- if self.recon_target == "tomo":
- img_rec = np.reshape(step_result['recon'], (self.img_w, self.img_w))
- prj_rec = np.reshape(step_result['prj_rec'], (self.img_h, self.img_w))
- elif self.recon_target == "tensor":
- img_rec = np.reshape(step_result['recon'][:,:,:,0], (self.img_w, self.img_w))
- prj_rec = np.reshape(step_result['prj_rec'], (self.img_h, self.img_w))
- elif self.recon_target == "phase":
- img_rec = np.reshape(step_result['phase'], (self.img_w, self.img_w))
- prj_rec = np.reshape(step_result['i_rec'], (self.img_h, self.img_w))
-
-
- # img_diff = np.abs(prj_rec - self.img_input)
- img_range = self.img_input.max() - self.img_input.min()
- ssim_index, ssim_map = ssim(prj_rec, self.img_input, full=True, data_range=img_range)
- self.tx1.set_text("SSIM map of " + self.plot_txt + " for iteration {0}".format(self.epoch))
- vmax = np.max(ssim_map)
- vmin = np.min(ssim_map)
- self.im1.set_data(ssim_map)
- self.im1.set_clim(vmin, vmax)
- self.im2.set_data(img_rec)
- vmax = np.max(img_rec)
- vmin = np.min(img_rec)
- self.im2.set_clim(vmin, vmax)
- self.axs[1, 1].plot(self.plot_x, self.plot_y1, "r-")
- self.axs[1, 1].plot(self.plot_x, self.plot_y2, "b-")
-
- plt.tight_layout()
- if in_notebook():
- clear_output(wait=True)
- display(self.fig)
- plt.pause(0.001)
+ self.update_rate = int(update_rate)
+ self.step = 0
+ self.total_steps = 0 # set by caller for progress display
+ self._in_notebook = in_notebook()
+
+ # Accumulated loss history for inline display
+ self._g_losses = []
+ self._d_losses = []
+
+ if logdir is None:
+ ts = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
+ logdir = os.path.join("logs", "ganrec", ts)
+ self.logdir = logdir
+ self.writer = tf.summary.create_file_writer(logdir)
+
+ # Map recon_target to the dict keys produced by forward_fn
+ if recon_target in ("tomo", "tensor"):
+ self._recon_key = "recon"
+ self._pred_key = "prj_rec"
+ elif recon_target == "phase":
+ self._recon_key = "phase"
+ self._pred_key = "i_rec"
+ else:
+ self._recon_key = "recon"
+ self._pred_key = "predicted"
+
+ # Also accept "predicted" as a fallback for _pred_key
+ self._pred_fallback = "predicted"
+
+ # Write the input image once
+ with self.writer.as_default():
+ inp_img = self._to_image_tensor(
+ tf.constant(self.img_input))
+ tf.summary.image("input", inp_img, step=0)
+ self.writer.flush()
+
+ # ------------------------------------------------------------------ #
+ # Helpers
+ # ------------------------------------------------------------------ #
+ @staticmethod
+ def _to_image_tensor(t):
+ """Normalise a 2-D+ tensor to [0, 1] and reshape to [1, H, W, 1]."""
+ t = tf.cast(t, tf.float32)
+ t_flat = tf.reshape(t, [-1])
+ t_min = tf.reduce_min(t_flat)
+ t_max = tf.reduce_max(t_flat)
+ t_norm = (t - t_min) / (t_max - t_min + 1e-8)
+ if t_norm.shape.ndims is None or t_norm.shape.ndims == 2:
+ return tf.reshape(t_norm, [1, tf.shape(t_norm)[0], tf.shape(t_norm)[1], 1])
+ elif t_norm.shape.ndims == 3:
+ return tf.reshape(t_norm, [1, tf.shape(t_norm)[0], tf.shape(t_norm)[1], tf.shape(t_norm)[2]])
+ elif t_norm.shape.ndims == 4:
+ return t_norm[:1]
+ return tf.reshape(t_norm, [1, tf.shape(t_norm)[0], tf.shape(t_norm)[1], 1])
+
+ @staticmethod
+ def _to_png_bytes(img_4d):
+ """Convert a [1, H, W, 1] float32 tensor to PNG bytes."""
+ img_uint8 = tf.cast(tf.clip_by_value(img_4d[0] * 255.0, 0, 255), tf.uint8)
+ return tf.io.encode_png(img_uint8)
+
+ @tf.function(reduce_retracing=True)
+ def _tf_ssim(self, pred, target):
+ """Compute SSIM between two 2-D tensors (on device)."""
+ pred_4d = tf.reshape(pred, [1, tf.shape(pred)[0], tf.shape(pred)[1], 1])
+ tgt_4d = tf.reshape(target, [1, tf.shape(target)[0], tf.shape(target)[1], 1])
+ max_val = tf.reduce_max(tgt_4d) - tf.reduce_min(tgt_4d)
+ return tf.image.ssim(pred_4d, tgt_4d, max_val=max_val + 1e-8)
+
+ def _get_pred_tensor(self, step_result):
+ """Resolve the predicted-measurement tensor from step_result."""
+ if self._pred_key in step_result:
+ return step_result[self._pred_key]
+ return step_result.get(self._pred_fallback)
+
+ # ------------------------------------------------------------------ #
+ # Inline Jupyter display (no matplotlib)
+ # ------------------------------------------------------------------ #
+ @staticmethod
+ def _colorbar_html(tensor, height=256):
+ """Generate a vertical HTML/CSS colorbar with min/max value labels.
+
+ Uses a grayscale gradient matching the image normalization
+ (white = max, black = min) and annotates with true data values.
+ """
+ if hasattr(tensor, 'numpy'):
+ t_np = tensor.numpy()
+ else:
+ t_np = np.asarray(tensor)
+ vmin, vmax = float(t_np.min()), float(t_np.max())
+ return (
+ f''
+ f'
{vmax:.3g}
'
+ f'
'
+ f'
{vmin:.3g}
'
+ f'
'
+ f'
'
+ f'
{title}
'
+ f'
'
+ f'
'
+ f'{cbar}'
+ f'
'
+ f'Step {self.step}/{self.total_steps} ({pct:.0f}%)
'
+ )
+
+ parts.append(
+ f"G_loss: {self._g_losses[-1]:.4f} "
+ f"D_loss: {self._d_losses[-1]:.4f} "
+ f"SSIM: {ssim_val:.4f}"
+ )
+
+ # Row 1: Reconstruction
+ if self._recon_key in step_result:
+ recon_t = step_result[self._recon_key]
+ if self.recon_target == "tensor":
+ recon_t = recon_t[:, :, :, 0] if recon_t.shape.ndims == 4 else recon_t
+ rec_img = self._to_image_tensor(recon_t)
+ rec_png = self._to_png_bytes(rec_img).numpy()
+ rec_b64 = base64.b64encode(rec_png).decode()
+ parts.append(
+ ''
+ )
+ parts.append(self._image_panel("Reconstruction", recon_t, rec_b64))
+ parts.append("
")
+
+ # Row 2: Input + Predicted
+ parts.append(
+ ''
+ )
+ inp_t = tf.constant(self.img_input)
+ inp_img = self._to_image_tensor(inp_t)
+ inp_png = self._to_png_bytes(inp_img).numpy()
+ inp_b64 = base64.b64encode(inp_png).decode()
+ parts.append(self._image_panel("Input", inp_t, inp_b64))
+
+ pred_t = self._get_pred_tensor(step_result)
+ if pred_t is not None:
+ pred_img = self._to_image_tensor(pred_t)
+ pred_png = self._to_png_bytes(pred_img).numpy()
+ pred_b64 = base64.b64encode(pred_png).decode()
+ parts.append(self._image_panel("Predicted", pred_t, pred_b64))
+ parts.append("
")
+
+ # Row 3: Loss convergence curve
+ if len(self._g_losses) > 1:
+ parts.append(self._loss_svg())
+
+ clear_output(wait=True)
+ display(HTML("\n".join(parts)))
+
+ def _loss_svg(self):
+ """Render a compact inline SVG loss chart with axes (no matplotlib)."""
+ w, h = 600, 160
+ left, right, top, bottom = 60, 20, 15, 40
+ plot_w = w - left - right
+ plot_h = h - top - bottom
+ n = len(self._g_losses)
+ if n < 2:
+ return ""
+
+ # Combine both loss series to get a shared Y range
+ all_vals = self._g_losses + self._d_losses
+ ymin, ymax = min(all_vals), max(all_vals)
+ if ymax == ymin:
+ ymax = ymin + 1
+
+ def _map(vals):
+ xs = [left + i / (n - 1) * plot_w for i in range(n)]
+ ys = [top + (1 - (v - ymin) / (ymax - ymin)) * plot_h for v in vals]
+ return xs, ys
+
+ gx, gy = _map(self._g_losses)
+ dx, dy = _map(self._d_losses)
+
+ g_pts = " ".join(f"{x:.1f},{y:.1f}" for x, y in zip(gx, gy))
+ d_pts = " ".join(f"{x:.1f},{y:.1f}" for x, y in zip(dx, dy))
+
+ # X-axis tick labels (step numbers)
+ n_xticks = min(5, n)
+ xticks = ""
+ for i in range(n_xticks):
+ idx = int(i / (n_xticks - 1) * (n - 1)) if n_xticks > 1 else 0
+ x_pos = left + idx / (n - 1) * plot_w
+ step_num = idx # monitor step index
+ xticks += (
+ f''
+ f'{step_num}'
+ )
+
+ # Y-axis tick labels (loss values)
+ n_yticks = 4
+ yticks = ""
+ for i in range(n_yticks):
+ frac = i / (n_yticks - 1)
+ y_pos = top + (1 - frac) * plot_h
+ val = ymin + frac * (ymax - ymin)
+ yticks += (
+ f''
+ f'{val:.3g}'
+ )
+
+ svg = (
+ f''
+ )
+ return svg
+
+ # ------------------------------------------------------------------ #
+ # Public API
+ # ------------------------------------------------------------------ #
+ def update_plot(self, step_result):
+ """Write summaries and optionally display inline in Jupyter.
+
+ Parameters
+ ----------
+ step_result : dict
+ Dictionary returned by ``recon_step``.
+ """
+ self.step += 1
+ step = self.step
+
+ # Accumulate losses (cheap — just float scalars)
+ self._g_losses.append(float(step_result["g_loss"].numpy()))
+ self._d_losses.append(float(step_result["d_loss"].numpy()))
+
+ with self.writer.as_default():
+ # ---- Scalar summaries (every step) ----
+ g_loss = step_result["g_loss"]
+ d_loss = step_result["d_loss"]
+ tf.summary.scalar("loss/generator", g_loss, step=step)
+ tf.summary.scalar("loss/discriminator", d_loss, step=step)
+
+ if "g_gnorm" in step_result:
+ tf.summary.scalar("grad_norm/generator",
+ step_result["g_gnorm"], step=step)
+ if "d_gnorm" in step_result:
+ tf.summary.scalar("grad_norm/discriminator",
+ step_result["d_gnorm"], step=step)
+
+ # ---- Image summaries + SSIM (every update_rate steps) ----
+ ssim_val = 0.0
+ if step % self.update_rate == 0:
+ # Reconstruction image
+ if self._recon_key in step_result:
+ recon_t = step_result[self._recon_key]
+ if self.recon_target == "tensor":
+ recon_t = recon_t[:, :, :, 0] if recon_t.shape.ndims == 4 else recon_t
+ recon_img = self._to_image_tensor(recon_t)
+ tf.summary.image("reconstruction", recon_img, step=step)
+
+ # Predicted measurement + SSIM
+ pred_t = self._get_pred_tensor(step_result)
+ if pred_t is not None:
+ pred_img = self._to_image_tensor(pred_t)
+ tf.summary.image("predicted", pred_img, step=step)
+
+ pred_2d = tf.reshape(
+ tf.cast(pred_t, tf.float32),
+ [self.img_h, self.img_w],
+ )
+ input_tf = tf.constant(self.img_input)
+ ssim_val = float(self._tf_ssim(pred_2d, input_tf)[0].numpy())
+ tf.summary.scalar("metrics/ssim", ssim_val, step=step)
+
+ # Flush + inline display periodically
+ if step % self.update_rate == 0:
+ self.writer.flush()
+ if self._in_notebook:
+ self._display_inline(step_result, ssim_val)
+
+ def close_plot(self, last_step_result=None):
+ """Flush, do a final inline display, and close the TensorBoard writer."""
+ self.writer.flush()
+ # Final inline display so progress shows 100%
+ if self._in_notebook and last_step_result is not None:
+ # Compute final SSIM
+ ssim_val = 0.0
+ pred_t = self._get_pred_tensor(last_step_result)
+ if pred_t is not None:
+ pred_2d = tf.reshape(
+ tf.cast(pred_t, tf.float32),
+ [self.img_h, self.img_w],
+ )
+ input_tf = tf.constant(self.img_input)
+ ssim_val = float(self._tf_ssim(pred_2d, input_tf)[0].numpy())
+ self._display_inline(last_step_result, ssim_val)
+ self.writer.close()
- def close_plot(self):
- plt.close()
def display_strain_tensor(tensor, profile_index=None):
"""
@@ -243,6 +508,8 @@ def display_strain_tensor(tensor, profile_index=None):
if tensor.shape[0] not in [3, 6]:
raise ValueError("Input tensor must have 3 or 6 components in the first dimension")
+ import matplotlib.pyplot as plt
+
component_names = (
[
r"$\epsilon_{xx}$",
diff --git a/ganrectorch/config.json b/ganrectorch/config.json
index 793cc82..8d56139 100644
--- a/ganrectorch/config.json
+++ b/ganrectorch/config.json
@@ -6,7 +6,7 @@
"dropout": 0.25,
"l1_ratio": 100,
"g_learning_rate": 1e-3,
- "d_learning_rate": 1e-3,
+ "d_learning_rate": 1e-5,
"save_wpath": null,
"init_wpath": null,
"init_model": false,
diff --git a/ganrectorch/ganrec.py b/ganrectorch/ganrec.py
index 050410e..f694985 100644
--- a/ganrectorch/ganrec.py
+++ b/ganrectorch/ganrec.py
@@ -1,6 +1,8 @@
import os
import json
+import copy
from tqdm import tqdm
+import numpy as np
import torch
import torch.nn.functional as F
from torch import nn, optim
@@ -33,6 +35,64 @@ def load_config(filename):
config = load_config("config.json")
+class DeviceEMA:
+ """Exponential moving average of model parameters (on device)."""
+ def __init__(self, model, decay=0.999):
+ self.model = model
+ self.decay = decay
+ self.shadow = {name: p.data.clone() for name, p in model.named_parameters()}
+ self.backup = {}
+
+ @torch.no_grad()
+ def update(self):
+ for name, p in self.model.named_parameters():
+ self.shadow[name].mul_(self.decay).add_(p.data, alpha=1.0 - self.decay)
+
+ def swap_in(self):
+ """Load EMA weights into the model (backup current weights)."""
+ self.backup = {name: p.data.clone() for name, p in self.model.named_parameters()}
+ for name, p in self.model.named_parameters():
+ p.data.copy_(self.shadow[name])
+
+ def swap_out(self):
+ """Restore original (non-EMA) weights."""
+ for name, p in self.model.named_parameters():
+ p.data.copy_(self.backup[name])
+ self.backup = {}
+
+
+class DeviceSnapshot:
+ """On-device snapshot/restore of model state dicts."""
+ def __init__(self, models):
+ self.models = models
+ self.copies = [copy.deepcopy(m.state_dict()) for m in models]
+
+ @torch.no_grad()
+ def snapshot(self):
+ for i, m in enumerate(self.models):
+ self.copies[i] = copy.deepcopy(m.state_dict())
+
+ @torch.no_grad()
+ def restore(self):
+ for m, sd in zip(self.models, self.copies):
+ m.load_state_dict(sd)
+
+
+class ScalarEMA:
+ """Exponential moving average of a scalar value."""
+ def __init__(self, decay=0.95):
+ self.decay = decay
+ self.value = None
+
+ def update(self, x):
+ x = float(x)
+ if self.value is None:
+ self.value = x
+ else:
+ self.value = self.decay * self.value + (1.0 - self.decay) * x
+ return self.value
+
+
# @torch.compile()
def discriminator_loss(real_output, fake_output):
real_loss = torch.mean(torch.nn.BCEWithLogitsLoss()(real_output, torch.ones_like(real_output)))
@@ -119,8 +179,14 @@ def make_model(self):
self.prj_input.shape[2], self.prj_input.shape[3], self.conv_num, self.conv_size, self.dropout, 1
)
self.discriminator = Discriminator()
- self.generator_optimizer = optim.Adam(self.generator.parameters(), lr=self.g_learning_rate)
- self.discriminator_optimizer = optim.Adam(self.discriminator.parameters(), lr=self.d_learning_rate)
+ self.generator_optimizer = optim.AdamW(
+ self.generator.parameters(), lr=self.g_learning_rate,
+ weight_decay=1e-4, betas=(0.9, 0.99)
+ )
+ self.discriminator_optimizer = optim.AdamW(
+ self.discriminator.parameters(), lr=self.d_learning_rate,
+ weight_decay=1e-4, betas=(0.9, 0.99)
+ )
# @torch.compile()
def nor_tomo(self, data):
@@ -140,7 +206,6 @@ def nor_tomo(self, data):
return shifted_data
- @torch.compile()
def recon_step(self, prj, ang):
self.generator_optimizer.zero_grad()
self.discriminator_optimizer.zero_grad()
@@ -149,8 +214,6 @@ def recon_step(self, prj, ang):
recon = self.nor_tomo(recon)
prj_rec = self.radon(recon, ang)
prj_rec = self.nor_tomo(prj_rec)
- # print(f"prj shape: {prj.shape}")
- # print(f"prj_rec shape: {prj_rec.shape}")
real_output = self.discriminator(prj)
fake_output = self.discriminator(prj_rec)
g_loss = generator_loss(fake_output, prj, prj_rec, self.l1_ratio)
@@ -160,6 +223,14 @@ def recon_step(self, prj, ang):
self.scaler.scale(g_loss).backward(retain_graph=True)
self.scaler.scale(d_loss).backward()
+ # Unscale before clipping
+ self.scaler.unscale_(self.generator_optimizer)
+ self.scaler.unscale_(self.discriminator_optimizer)
+
+ # Gradient clipping for stability
+ torch.nn.utils.clip_grad_norm_(self.generator.parameters(), max_norm=1.0)
+ torch.nn.utils.clip_grad_norm_(self.discriminator.parameters(), max_norm=1.0)
+
# Optimizer step with gradient scaling
self.scaler.step(self.generator_optimizer)
self.scaler.step(self.discriminator_optimizer)
@@ -178,43 +249,404 @@ def recon(self):
self.generator.load_state_dict(torch.load(self.init_wpath + "generator.pth"))
print("generator is initialized")
self.discriminator.load_state_dict(torch.load(self.init_wpath + "discriminator.pth"))
- recon = torch.zeros((self.iter_num, 1, self.px, self.px))
- gen_loss = torch.zeros((self.iter_num))
- ###########################################################################
- # Reconstruction process monitor
+ # ---------- Stability tunables ----------
+ ema_decay = float(getattr(self, "ema_decay", 0.99))
+ snapshot_every = int(getattr(self, "snapshot_every", 50))
+ log_every = int(getattr(self, "log_every", 10))
+ spike_factor = float(getattr(self, "spike_factor", 1.5))
+ warmup_steps = int(getattr(self, "warmup_steps", max(5, self.iter_num // 20)))
+ lr_backoff = float(getattr(self, "lr_backoff", 0.5))
+ lr_floor = float(getattr(self, "lr_floor", 1e-6))
+ freeze_disc_max = int(getattr(self, "freeze_disc_max", 10))
+
+ # ---------- EMA & Snapshots ----------
+ ema = DeviceEMA(self.generator, decay=ema_decay)
+ snap = DeviceSnapshot([self.generator, self.discriminator])
+ snap.snapshot()
+
+ g_ema = ScalarEMA(0.95)
+ d_ema = ScalarEMA(0.95)
+
+ g_lr0 = self.g_learning_rate
+ d_lr0 = self.d_learning_rate
+ freeze_disc_steps = 0
+
+ # ---------- Monitor ----------
+ gen_loss = torch.zeros(self.iter_num)
if self.recon_monitor:
plot_x, plot_loss = [], []
recon_monitor = RECONmonitor("tomo", self.prj_input.cpu())
- pbar = tqdm(total=self.iter_num, desc="Reconstruction Progress", position=0, leave=True)
+ pbar = tqdm(total=self.iter_num, desc="Reconstruction Progress", position=0, leave=True)
- ###########################################################################
+ recon = None
for epoch in range(self.iter_num):
-
- ###########################################################################
- ## Call the rconstruction step
-
+ # Apply temporary D freeze by setting LR to 0
+ if freeze_disc_steps > 0:
+ for pg in self.discriminator_optimizer.param_groups:
+ pg["lr"] = 0.0
+ freeze_disc_steps -= 1
+ else:
+ for pg in self.discriminator_optimizer.param_groups:
+ pg["lr"] = d_lr0
+
+ # ---------- Reconstruction step ----------
step_result = self.recon_step(self.prj_input, self.angle)
+ g_loss_val = step_result["g_loss"].item()
+ d_loss_val = step_result["d_loss"].item()
recon = step_result["recon"]
- gen_loss[epoch] = step_result["g_loss"]
- ###########################################################################
+
+ # ---------- NaN / Inf guard ----------
+ if not (np.isfinite(g_loss_val) and np.isfinite(d_loss_val)):
+ snap.restore()
+ # Backoff learning rates
+ for pg in self.generator_optimizer.param_groups:
+ pg["lr"] = max(lr_floor, pg["lr"] * lr_backoff)
+ for pg in self.discriminator_optimizer.param_groups:
+ pg["lr"] = max(lr_floor, pg["lr"] * lr_backoff)
+ freeze_disc_steps = freeze_disc_max
+ pbar.set_postfix_str("NaN/Inf→rollback")
+ pbar.update(1)
+ continue
+
+ # ---------- Flat output guard ----------
+ if epoch % log_every == 0:
+ recon_var = torch.var(recon).item()
+ if recon_var < 1e-4:
+ snap.restore()
+ for pg in self.discriminator_optimizer.param_groups:
+ pg["lr"] = max(lr_floor, pg["lr"] * lr_backoff)
+ pbar.set_postfix_str("flat→rollback")
+ pbar.update(1)
+ continue
+
+ # ---------- Update scalar EMAs ----------
+ g_bar = g_ema.update(g_loss_val)
+ d_bar = d_ema.update(d_loss_val)
+
+ # ---------- Spike detection ----------
+ if epoch > warmup_steps and epoch % log_every == 0:
+ if (g_loss_val > spike_factor * max(1e-8, g_bar) or
+ d_loss_val > spike_factor * max(1e-8, d_bar)):
+ snap.restore()
+ for pg in self.generator_optimizer.param_groups:
+ pg["lr"] = max(lr_floor, pg["lr"] * lr_backoff)
+ for pg in self.discriminator_optimizer.param_groups:
+ pg["lr"] = max(lr_floor, pg["lr"] * lr_backoff)
+ freeze_disc_steps = freeze_disc_max
+ pbar.set_postfix_str("spike→rollback")
+ pbar.update(1)
+ continue
+
+ # ---------- Good step: update EMA and snapshot ----------
+ ema.update()
+ if epoch % snapshot_every == 0:
+ snap.snapshot()
+
+ gen_loss[epoch] = g_loss_val
+
+ # ---------- Logging / monitor ----------
if self.recon_monitor:
plot_x.append(epoch)
plot_loss = gen_loss[: epoch + 1]
- pbar.set_postfix(G_loss=gen_loss[epoch].item(), D_loss=step_result["d_loss"].item())
- pbar.update(1)
+ pbar.set_postfix(G_loss=f"{g_loss_val:.4f}", D_loss=f"{d_loss_val:.4f}")
+ pbar.update(1)
+
if (epoch + 1) % 100 == 0:
if self.recon_monitor:
prj_rec = step_result["prj_rec"].view(self.nang, self.px)
prj_diff = torch.abs(prj_rec - self.prj_input.view((self.nang, self.px))).cpu()
rec_plt = recon.view(self.px, self.px).cpu()
recon_monitor.update_plot(epoch, prj_diff, rec_plt, plot_x, plot_loss.cpu())
- # print('Iteration {}: G_loss is {} and D_loss is {}'.format(epoch + 1,
- # gen_loss[epoch],
- # step_result['d_loss'].item()))
- if self.save_wpath != None:
+
+ pbar.close()
+
+ if self.save_wpath is not None:
torch.save(self.generator.state_dict(), self.save_wpath + "generator.pth")
torch.save(self.discriminator.state_dict(), self.save_wpath + "discriminator.pth")
if self.recon_monitor:
recon_monitor.close_plot()
- return tensor_to_np(recon.cpu())
+
+ # Final output: use EMA weights for the result
+ ema.swap_in()
+ try:
+ with torch.no_grad():
+ final_recon = self.generator(self.prj_input)
+ final_recon = self.nor_tomo(final_recon)
+ finally:
+ ema.swap_out()
+
+ return tensor_to_np(final_recon.cpu())
+
+
+class GANrec:
+ """
+ General-purpose GAN-based reconstruction with full stability features.
+
+ Unlike specialized classes (GANtomo, etc.), GANrec is not restricted
+ to a specific forward model. The user provides a ``forward_fn`` callable
+ that defines the physics mapping from reconstruction to measurement space.
+
+ Stability features:
+ EMA of generator weights, on-device rollback snapshots, NaN/Inf guards,
+ loss-spike detection, discriminator freezing, gradient clipping, and
+ learning-rate backoff.
+
+ Parameters
+ ----------
+ input_data : ndarray (2-D)
+ Measured data (sinogram, intensity image, diffraction pattern, …).
+ forward_fn : callable
+ ``(gen_output, input_tensor) -> dict``
+ Must return a dict with at least ``"predicted"`` (simulated measurement,
+ same shape as *input_tensor*). May include extra keys (``"recon"``,
+ ``"phase"``, …) carried through for monitoring / output extraction.
+ output_num : int, optional
+ Number of generator output channels (default 1).
+ output_key : str, optional
+ Key in ``forward_fn`` result for the final output (default ``"recon"``).
+ shape_output : tuple, optional
+ Output reshape target. Defaults to ``(input_data.shape[-1],) * 2``.
+ **kwargs
+ Override config values (``iter_num``, ``l1_ratio``, ``g_learning_rate``,
+ ``d_learning_rate``, ``conv_num``, …).
+
+ Examples
+ --------
+ **Tomography**::
+
+ from ganrectorch.propagators import RadonTransform
+
+ radon = RadonTransform(torch.empty(1, 1, px, px), angle)
+ radon = radon.to(device)
+
+ def tomo_forward(gen_output, input_tensor):
+ recon = nor_tomo(gen_output)
+ prj_rec = radon(recon, angle)
+ prj_rec = nor_tomo(prj_rec)
+ return {"recon": recon, "predicted": prj_rec}
+
+ gan = GANrec(sinogram, tomo_forward)
+ result = gan.recon()
+
+ **Phase retrieval (Fresnel)**::
+
+ from ganrectorch.propagators import PhaseFresnel
+
+ def phase_forward(gen_output, input_tensor):
+ phase = normalize_phase(gen_output[:, 0:1, :, :])
+ absorption = normalize_phase(gen_output[:, 1:2, :, :]) * abs_ratio
+ i_rec = PhaseFresnel(phase[0,0], absorption[0,0], ff, px).compute()
+ return {"phase": phase, "absorption": absorption, "predicted": i_rec}
+
+ gan = GANrec(intensity, phase_forward, output_num=2, output_key="phase")
+ result = gan.recon()
+ """
+
+ def __init__(self, input_data, forward_fn, output_num=1, output_key="recon",
+ shape_output=None, **kwargs):
+ base_args = config["GANtomo"].copy()
+ base_args.update(kwargs)
+ super().__init__()
+ torch_configures()
+ self.scaler = GradScaler()
+ self.forward_fn = forward_fn
+ self.output_num = output_num
+ self.output_key = output_key
+
+ # Store and prepare input
+ self.input_data = input_data
+ self.input_shape = input_data.shape
+ self.shape_output = shape_output or (self.input_shape[-1], self.input_shape[-1])
+ self.input_tensor = (
+ torch.from_numpy(input_data).float().unsqueeze(0).unsqueeze(0)
+ ) # [1, 1, H, W]
+
+ self.iter_num = base_args["iter_num"]
+ self.conv_num = base_args["conv_num"]
+ self.conv_size = base_args["conv_size"]
+ self.dropout = base_args["dropout"]
+ self.l1_ratio = base_args["l1_ratio"]
+ self.g_learning_rate = base_args["g_learning_rate"]
+ self.d_learning_rate = base_args["d_learning_rate"]
+ self.save_wpath = base_args["save_wpath"]
+ self.init_wpath = base_args["init_wpath"]
+ self.init_model = base_args["init_model"]
+ self.recon_monitor = base_args["recon_monitor"]
+ self.generator = None
+ self.discriminator = None
+
+ def make_model(self):
+ h, w = self.input_shape
+ self.generator = Generator(h, w, self.conv_num, self.conv_size,
+ self.dropout, self.output_num)
+ self.discriminator = Discriminator()
+ self.generator_optimizer = optim.AdamW(
+ self.generator.parameters(), lr=self.g_learning_rate,
+ weight_decay=1e-4, betas=(0.9, 0.99),
+ )
+ self.discriminator_optimizer = optim.AdamW(
+ self.discriminator.parameters(), lr=self.d_learning_rate,
+ weight_decay=1e-4, betas=(0.9, 0.99),
+ )
+
+ def recon_step(self, input_tensor):
+ self.generator_optimizer.zero_grad()
+ self.discriminator_optimizer.zero_grad()
+
+ with autocast():
+ gen_output = self.generator(input_tensor)
+ fwd = self.forward_fn(gen_output, input_tensor)
+ predicted = fwd["predicted"]
+ real_output = self.discriminator(input_tensor)
+ fake_output = self.discriminator(predicted)
+ g_loss = generator_loss(fake_output, input_tensor, predicted,
+ self.l1_ratio)
+ d_loss = discriminator_loss(real_output, fake_output)
+
+ self.scaler.scale(g_loss).backward(retain_graph=True)
+ self.scaler.scale(d_loss).backward()
+ self.scaler.unscale_(self.generator_optimizer)
+ self.scaler.unscale_(self.discriminator_optimizer)
+ torch.nn.utils.clip_grad_norm_(self.generator.parameters(), max_norm=1.0)
+ torch.nn.utils.clip_grad_norm_(self.discriminator.parameters(), max_norm=1.0)
+ self.scaler.step(self.generator_optimizer)
+ self.scaler.step(self.discriminator_optimizer)
+ self.scaler.update()
+
+ fwd["g_loss"] = g_loss
+ fwd["d_loss"] = d_loss
+ return fwd
+
+ def recon(self):
+ """Run the full reconstruction loop with stability safeguards.
+
+ Returns
+ -------
+ ndarray
+ Reconstruction reshaped to ``shape_output``.
+ """
+ self.make_model()
+ self.input_tensor, self.generator, self.discriminator = to_device(
+ [self.input_tensor, self.generator, self.discriminator]
+ )
+
+ if self.init_wpath:
+ self.generator.load_state_dict(
+ torch.load(self.init_wpath + "generator.pth"))
+ self.discriminator.load_state_dict(
+ torch.load(self.init_wpath + "discriminator.pth"))
+ print("Models are initialized")
+
+ # ---------- Stability tunables ----------
+ ema_decay = float(getattr(self, "ema_decay", 0.99))
+ snapshot_every = int(getattr(self, "snapshot_every", 50))
+ log_every = int(getattr(self, "log_every", 10))
+ spike_factor = float(getattr(self, "spike_factor", 1.5))
+ warmup_steps = int(getattr(self, "warmup_steps",
+ max(5, self.iter_num // 20)))
+ lr_backoff = float(getattr(self, "lr_backoff", 0.5))
+ lr_floor = float(getattr(self, "lr_floor", 1e-6))
+ freeze_disc_max = int(getattr(self, "freeze_disc_max", 10))
+
+ # ---------- EMA & Snapshots ----------
+ ema = DeviceEMA(self.generator, decay=ema_decay)
+ snap = DeviceSnapshot([self.generator, self.discriminator])
+ snap.snapshot()
+
+ g_ema_s = ScalarEMA(0.95)
+ d_ema_s = ScalarEMA(0.95)
+ d_lr0 = self.d_learning_rate
+ freeze_disc_steps = 0
+
+ # ---------- Monitor ----------
+ pbar = tqdm(total=self.iter_num, desc="Reconstruction", leave=True)
+
+ step_result = {}
+ for step in range(self.iter_num):
+ # D freeze management
+ if freeze_disc_steps > 0:
+ for pg in self.discriminator_optimizer.param_groups:
+ pg["lr"] = 0.0
+ freeze_disc_steps -= 1
+ else:
+ for pg in self.discriminator_optimizer.param_groups:
+ pg["lr"] = d_lr0
+
+ # ---- step ----
+ step_result = self.recon_step(self.input_tensor)
+ g_loss_val = step_result["g_loss"].item()
+ d_loss_val = step_result["d_loss"].item()
+
+ # NaN / Inf guard
+ if not (np.isfinite(g_loss_val) and np.isfinite(d_loss_val)):
+ snap.restore()
+ for pg in self.generator_optimizer.param_groups:
+ pg["lr"] = max(lr_floor, pg["lr"] * lr_backoff)
+ for pg in self.discriminator_optimizer.param_groups:
+ pg["lr"] = max(lr_floor, pg["lr"] * lr_backoff)
+ freeze_disc_steps = freeze_disc_max
+ pbar.set_postfix_str("NaN/Inf→rollback")
+ pbar.update(1)
+ continue
+
+ # Flat-output guard
+ if self.output_key in step_result and step % log_every == 0:
+ if torch.var(step_result[self.output_key]).item() < 1e-4:
+ snap.restore()
+ for pg in self.discriminator_optimizer.param_groups:
+ pg["lr"] = max(lr_floor, pg["lr"] * lr_backoff)
+ pbar.set_postfix_str("flat→rollback")
+ pbar.update(1)
+ continue
+
+ # Update scalar EMAs
+ g_bar = g_ema_s.update(g_loss_val)
+ d_bar = d_ema_s.update(d_loss_val)
+
+ # Spike detection
+ if step > warmup_steps and step % log_every == 0:
+ if (g_loss_val > spike_factor * max(1e-8, g_bar) or
+ d_loss_val > spike_factor * max(1e-8, d_bar)):
+ snap.restore()
+ for pg in self.generator_optimizer.param_groups:
+ pg["lr"] = max(lr_floor, pg["lr"] * lr_backoff)
+ for pg in self.discriminator_optimizer.param_groups:
+ pg["lr"] = max(lr_floor, pg["lr"] * lr_backoff)
+ freeze_disc_steps = freeze_disc_max
+ pbar.set_postfix_str("spike→rollback")
+ pbar.update(1)
+ continue
+
+ # Good step
+ ema.update()
+ if step % snapshot_every == 0:
+ snap.snapshot()
+
+ if step % log_every == 0:
+ pbar.set_postfix(G_loss=f"{g_loss_val:.4f}",
+ D_loss=f"{d_loss_val:.4f}")
+ pbar.update(1)
+
+ pbar.close()
+
+ if self.save_wpath is not None:
+ torch.save(self.generator.state_dict(),
+ self.save_wpath + "generator.pth")
+ torch.save(self.discriminator.state_dict(),
+ self.save_wpath + "discriminator.pth")
+
+ # ---------- Final output with EMA weights ----------
+ ema.swap_in()
+ try:
+ with torch.no_grad():
+ gen_output = self.generator(self.input_tensor)
+ fwd = self.forward_fn(gen_output, self.input_tensor)
+ if self.output_key in fwd:
+ final = fwd[self.output_key]
+ else:
+ final = fwd["predicted"]
+ finally:
+ ema.swap_out()
+
+ return tensor_to_np(final.cpu()).reshape(self.shape_output)