remove deterministic flags for refactor control

Author: KatarinaYuan

src/data/dataset/__pycache__/dataset_core.cpython-39.pyc

@@ -1,8 +1,6 @@
 """Module `data/dataset_core.py`."""
 import ctypes
 import ctypes.util
-import hashlib
-import os
 from logging import Logger
 from typing import Any, Dict, List, Tuple
 
@@ -123,17 +121,8 @@ def collate_fn(dataset, batch):
 def mapping_cells(dataset, ds_name, local_idx):
     """Execute `mapping_cells` and return values used by downstream logic."""
     assert dataset.data_args.mapping_strategy == "random"
-    deterministic_mapping = os.getenv("PDIFF_DETERMINISTIC_MAPPING", "0") == "1"
-    if deterministic_mapping:
-        base_seed = int(os.getenv("PDIFF_DETERMINISTIC_MAPPING_SEED", "20260304"))
-        key = f"{base_seed}|{ds_name}|{int(local_idx)}"
-        digest = hashlib.blake2b(key.encode("utf-8"), digest_size=8).digest()
-        rng = np.random.default_rng(int.from_bytes(digest, "little", signed=False))
-        randint = lambda high: int(rng.integers(0, high))
-        choice = lambda arr: arr[int(rng.integers(0, len(arr)))]
-    else:
-        randint = lambda high: int(np.random.randint(0, high, size=(1))[0])
-        choice = lambda arr: np.random.choice(arr, 1)[0]
+    randint = lambda high: int(np.random.randint(0, high, size=(1))[0])
+    choice = lambda arr: np.random.choice(arr, 1)[0]
 
     cache = dataset.meta_cache._cache[dataset.dataset_path_map[ds_name]]
 

src/data/dataset/dataset_core.py

@@ -1,8 +1,5 @@
 """Training/loss methods for GaussianDiffusion."""
 
-import json
-import os
-
 import torch as th
 from geomloss import SamplesLoss
 
@@ -43,23 +40,6 @@ def _build_training_target(model_mean_type, x_start, noise):
         ModelMeanType.EPSILON: noise,
     }[model_mean_type]
 
-
-
-def _safe_stats(x):
-    """Execute `_safe_stats` and return values used by downstream logic."""
-    if x is None:
-        return None
-    if not isinstance(x, th.Tensor):
-        return str(type(x))
-    y = x.detach().float()
-    return {
-        "shape": list(y.shape),
-        "mean": float(y.mean().item()),
-        "std": float(y.std(unbiased=False).item()),
-        "norm": float(y.norm().item()),
-    }
-
-
 class GaussianDiffusionTrainingMixin:
     """Gaussiandiffusiontrainingmixin implementation used by the PerturbDiff pipeline."""
     def get_model_output(
@@ -209,15 +189,9 @@ def training_losses(
 
         x_0 = th.zeros_like(x_t)
         control_0 = th.zeros_like(control_input_t)
-        selfcond_used = False
-        force_selfcond = os.getenv("PDIFF_FORCE_SELFCOND")
-        if force_selfcond is None:
-            use_selfcond_now = bool((th.rand(1) > 0.5).item())
-        else:
-            use_selfcond_now = force_selfcond.strip() in {"1", "true", "True"}
+        use_selfcond_now = bool((th.rand(1) > 0.5).item())
 
         if use_selfcond_now:
-            selfcond_used = True
             with th.no_grad():
                 out = self.get_model_output(
                     model=model,
@@ -230,7 +204,7 @@ def training_losses(
             x_0 = out["x"]
             control_0 = th.zeros_like(out["x_control"])
 
-        terms, model_output = self.diffusion_loss(
+        terms = self.diffusion_loss(
             model=model,
             x_start=x_start,
             x_t=x_t,
@@ -242,35 +216,10 @@ def training_losses(
             x_0=x_0,
             control_0=control_0,
             MMD_loss_fn=MMD_loss_fn,
-            return_model_output=True,
         )
 
         if model.model_cfg.no_mse_loss:
             terms["mse1"] = th.zeros_like(terms["mse1"])
-        trace_path = os.getenv("PDIFF_TRACE_PATH")
-        if trace_path and not getattr(self, "_pdiff_trace_dumped", False):
-            trace = {
-                "t_head": [int(v) for v in t.detach().cpu().view(-1)[:8]],
-                "x_start": _safe_stats(x_start),
-                "control_input_start": _safe_stats(control_input_start),
-                "noise": _safe_stats(noise),
-                "x_t": _safe_stats(x_t),
-                "control_input_t": _safe_stats(control_input_t),
-                "batch_emb_is_none": self_condition.get("batch_emb") is None,
-                "cont_emb_is_none": self_condition.get("cont_emb") is None,
-                "model_output_x": _safe_stats(model_output.get("x")),
-                "model_output_x_control": _safe_stats(model_output.get("x_control")),
-                "selfcond_used": bool(selfcond_used),
-                "x_0": _safe_stats(x_0),
-                "control_0": _safe_stats(control_0),
-                "terms_loss1_mean": float(terms["loss1"].detach().float().mean().item()),
-                "terms_mse1_mean": float(terms["mse1"].detach().float().mean().item()),
-                "terms_mmd1_mean": float(terms["mmd1"].detach().float().mean().item()) if "mmd1" in terms else None,
-            }
-            os.makedirs(os.path.dirname(trace_path), exist_ok=True)
-            with open(trace_path, "w", encoding="utf-8") as fout:
-                json.dump(trace, fout, indent=2)
-            self._pdiff_trace_dumped = True
         return terms
 
 

src/models/diffusion/__pycache__/diffusion_training.cpython-39.pyc

@@ -1,7 +1,6 @@
 """Lightning training module split from lightning_module (logic-preserving)."""
 
 import gc
-import os
 import pickle
 import sys
 import time
@@ -231,22 +230,15 @@ def _compute_loss(self, batch):
                 "ds_name": batch["ds_name"],
                 }
 
-        align_rng = os.getenv("PDIFF_FORCE_ALIGNED_RNG", "0") == "1"
-        if align_rng:
-            step_seed = int(os.getenv("PDIFF_SEED_BASE", "12345")) + int(self.global_step)
-            np.random.seed(step_seed)
-            torch.manual_seed(step_seed)
-
         t, weights = self.schedule_sampler.sample(pert_emb.shape[0], device)
-        noise = torch.randn_like(pert_emb, dtype=torch.float64) if align_rng else None
 
         losses = self.diffusion.training_losses(
             self.model, 
             pert_emb, 
             t, 
             self_condition=cond, 
             model_kwargs=None, 
-            noise=noise,
+            noise=None,
             p_drop_cond=self.model_cfg.p_drop_cond,
             MMD_loss_fn=self.loss_fn,
         )