My_python_job/utils update

thuml · bojunZhang-heng · Jul 14, 2025 · Jul 14, 2025 · Jul 14, 2025 · Jul 14, 2025
commit 520290d1394c082922d6ae289308132d180a78d7
diff --git a/PDE-Solving-StandardBenchmark/My_python_job/utils/normalizer.py b/PDE-Solving-StandardBenchmark/My_python_job/utils/normalizer.py
@@ -0,0 +1,114 @@
+import torch
+from tqdm import *
+
+
+class IdentityTransformer():
+    def __init__(self, X):
+        self.mean = X.mean(dim=0, keepdim=True)
+        self.std = X.std(dim=0, keepdim=True) + 1e-8
+
+    def to(self, device):
+        self.mean = self.mean.to(device)
+        self.std = self.std.to(device)
+        return self
+
+    def cuda(self):
+        self.mean = self.mean.cuda()
+        self.std = self.std.cuda()
+
+    def cpu(self):
+        self.mean = self.mean.cpu()
+        self.std = self.std.cpu()
+
+    def encode(self, x):
+        return x
+
+    def decode(self, x):
+        return x
+
+
+class UnitTransformer():
+    def __init__(self, X):
+        self.mean = X.mean(dim=(0, 1), keepdim=True)
+        self.std = X.std(dim=(0, 1), keepdim=True) + 1e-8
+
+    def to(self, device):
+        self.mean = self.mean.to(device)
+        self.std = self.std.to(device)
+        return self
+
+    def cuda(self):
+        self.mean = self.mean.cuda()
+        self.std = self.std.cuda()
+
+    def cpu(self):
+        self.mean = self.mean.cpu()
+        self.std = self.std.cpu()
+
+    def encode(self, x):
+        x = (x - self.mean) / (self.std)
+        return x
+
+    def decode(self, x):
+        return x * self.std + self.mean
+
+    def transform(self, X, inverse=True, component='all'):
+        if component == 'all' or 'all-reduce':
+            if inverse:
+                orig_shape = X.shape
+                return (X * (self.std - 1e-8) + self.mean).view(orig_shape)
+            else:
+                return (X - self.mean) / self.std
+        else:
+            if inverse:
+                orig_shape = X.shape
+                return (X * (self.std[:, component] - 1e-8) + self.mean[:, component]).view(orig_shape)
+            else:
+                return (X - self.mean[:, component]) / self.std[:, component]
+
+
+class UnitGaussianNormalizer(object):
+    def __init__(self, x, eps=0.00001, time_last=True):
+        super(UnitGaussianNormalizer, self).__init__()
+
+        self.mean = torch.mean(x, 0)
+        self.std = torch.std(x, 0)
+        self.eps = eps
+        self.time_last = time_last  # if the time dimension is the last dim
+
+    def encode(self, x):
+        x = (x - self.mean) / (self.std + self.eps)
+        return x
+
+    def decode(self, x, sample_idx=None):
+        # sample_idx is the spatial sampling mask
+        if sample_idx is None:
+            std = self.std + self.eps  # n
+            mean = self.mean
+        else:
+            if self.mean.ndim == sample_idx.ndim or self.time_last:
+                std = self.std[sample_idx] + self.eps  # batch*n
+                mean = self.mean[sample_idx]
+            if self.mean.ndim > sample_idx.ndim and not self.time_last:
+                std = self.std[..., sample_idx] + self.eps  # T*batch*n
+                mean = self.mean[..., sample_idx]
+        # x is in shape of batch*(spatial discretization size) or T*batch*(spatial discretization size)
+        x = (x * std) + mean
+        return x
+
+    def to(self, device):
+        if torch.is_tensor(self.mean):
+            self.mean = self.mean.to(device)
+            self.std = self.std.to(device)
+        else:
+            self.mean = torch.from_numpy(self.mean).to(device)
+            self.std = torch.from_numpy(self.std).to(device)
+        return self
+
+    def cuda(self):
+        self.mean = self.mean.cuda()
+        self.std = self.std.cuda()
+
+    def cpu(self):
+        self.mean = self.mean.cpu()
+        self.std = self.std.cpu()
diff --git a/PDE-Solving-StandardBenchmark/My_python_job/utils/testloss.py b/PDE-Solving-StandardBenchmark/My_python_job/utils/testloss.py
@@ -0,0 +1,45 @@
+import torch
+
+
+class TestLoss(object):
+    def __init__(self, d=2, p=2, size_average=True, reduction=True):
+        super(TestLoss, self).__init__()
+
+        assert d > 0 and p > 0
+
+        self.d = d
+        self.p = p
+        self.reduction = reduction
+        self.size_average = size_average
+
+    def abs(self, x, y):
+        num_examples = x.size()[0]
+
+        h = 1.0 / (x.size()[1] - 1.0)
+
+        all_norms = (h ** (self.d / self.p)) * torch.norm(x.view(num_examples, -1) - y.view(num_examples, -1), self.p,
+                                                          1)
+
+        if self.reduction:
+            if self.size_average:
+                return torch.mean(all_norms)
+            else:
+                return torch.sum(all_norms)
+
+        return all_norms
+
+    def rel(self, x, y):
+        num_examples = x.size()[0]
+
+        diff_norms = torch.norm(x.reshape(num_examples, -1) - y.reshape(num_examples, -1), self.p, 1)
+        y_norms = torch.norm(y.reshape(num_examples, -1), self.p, 1)
+        if self.reduction:
+            if self.size_average:
+                return torch.mean(diff_norms / y_norms)
+            else:
+                return torch.sum(diff_norms / y_norms)
+
+        return diff_norms / y_norms
+
+    def __call__(self, x, y):
+        return self.rel(x, y)