feat(ppsci): support data_effient_nopt for training and test

Author: xiaoyewww

examples/data_efficient_nopt/data_utils/datasets_ns.py

@@ -1214,7 +1214,7 @@
 #                 int((sigma * 4 + 1) / 2) * 2 + 1, (data_blur.shape[-1] // 2) * 2 - 1
 #             )
 #             if _kernel >= 2:
-#                 data_blur = functional.gaussian_blur(
+#                 data_blur = gaussian_blur(
 #                     data_blur, kernel_size=[_kernel, _kernel], sigma=sigma
 #                 )
 #             data_blur = data_blur.permute(1, 2, 0)

examples/data_efficient_nopt/models/fno.py

@@ -258,28 +258,31 @@ def forward_icl(self, x, demo_xs, demo_ys, use_tqdm=False):
         demo_xs_aug = []
         for _ in range(repeat):
             if sum(sigma_range) > 0:
-                raise NotImplementedError("Gaussian blur is not implemented in paddle.")
-                # sigma = random.uniform(*sigma_range)
-                # # https://github.com/scipy/scipy/blob/v1.11.4/scipy/ndimage/_filters.py#L232
-                # _kernel = min(
-                #     int((sigma * 4 + 1) / 2) * 2 + 1, (x.shape[1] // 2) * 2 - 1
-                # )
+                import random
+
+                sigma = random.uniform(*sigma_range)
+                # https://github.com/scipy/scipy/blob/v1.11.4/scipy/ndimage/_filters.py#L232
+                _kernel = min(
+                    int((sigma * 4 + 1) / 2) * 2 + 1, (x.shape[1] // 2) * 2 - 1
+                )
             mask = paddle.nn.functional.dropout(paddle.ones([1, C, H, W]), p=p)
             ######
+            from .gaussian_blur import gaussian_blur
+
             if sum(sigma_range) > 0:
-                # [TODO]
-                raise NotImplementedError("Gaussian blur is not implemented in paddle.")
-                # _x_aug = torchvision.transforms.functional.gaussian_blur(x.clone(), kernel_size=[_kernel, _kernel], sigma=sigma)
+                _x_aug = gaussian_blur(
+                    x.clone(), kernel_size=[_kernel, _kernel], sigma=sigma
+                )
             else:
                 _x_aug = x.clone()
             _x_aug = _x_aug * mask
             x_aug.append(_x_aug)
             ######
             _demo_xs_aug = []
             if sum(sigma_range) > 0:
-                # [TODO]
-                raise NotImplementedError("Gaussian blur is not implemented in paddle.")
-                # _demo_xs_aug = torchvision.transforms.functional.gaussian_blur(demo_xs.clone(), kernel_size=[_kernel, _kernel], sigma=sigma)
+                _demo_xs_aug = gaussian_blur(
+                    demo_xs.clone(), kernel_size=[_kernel, _kernel], sigma=sigma
+                )
             else:
                 _demo_xs_aug = demo_xs.clone()
             _demo_xs_aug = _demo_xs_aug * mask

examples/data_efficient_nopt/models/gaussian_blur.py

@@ -0,0 +1,105 @@
+from typing import List
+
+import paddle
+import paddle.tensor as Tensor
+
+
+def _cast_squeeze_in(img: Tensor, req_dtypes: List[paddle.dtype]):
+    need_squeeze = False
+    # make image NCHW
+    if img.ndim < 4:
+        img = img.unsqueeze(axis=0)
+        need_squeeze = True
+
+    out_dtype = img.dtype
+    need_cast = False
+    if out_dtype not in req_dtypes:
+        need_cast = True
+        req_dtype = req_dtypes[0]
+        img = img.to(req_dtype)
+    return img, need_cast, need_squeeze, out_dtype
+
+
+def _get_gaussian_kernel1d(
+    kernel_size: int, sigma: float, dtype: paddle.dtype
+) -> Tensor:
+    ksize_half = (kernel_size - 1) * 0.5
+
+    x = paddle.linspace(-ksize_half, ksize_half, num=kernel_size, dtype=dtype)
+    pdf = paddle.exp(-0.5 * (x / sigma).pow(2))
+    kernel1d = pdf / pdf.sum()
+
+    return kernel1d
+
+
+def _get_gaussian_kernel2d(
+    kernel_size: List[int], sigma: List[float], dtype: paddle.dtype
+) -> Tensor:
+    kernel1d_x = _get_gaussian_kernel1d(kernel_size[0], sigma[0], dtype)
+    kernel1d_y = _get_gaussian_kernel1d(kernel_size[1], sigma[1], dtype)
+    kernel2d = paddle.mm(kernel1d_y[:, None], kernel1d_x[None, :])
+    return kernel2d
+
+
+def _cast_squeeze_out(
+    img: Tensor, need_cast: bool, need_squeeze: bool, out_dtype: paddle.dtype
+) -> Tensor:
+    if need_squeeze:
+        img = img.squeeze(axis=0)
+
+    if need_cast:
+        if out_dtype in (
+            paddle.uint8,
+            paddle.int8,
+            paddle.int16,
+            paddle.int32,
+            paddle.int64,
+        ):
+            # it is better to round before cast
+            img = paddle.round(img)
+        img = img.to(out_dtype)
+
+    return img
+
+
+def gaussian_blur(img: Tensor, kernel_size: List[int], sigma: List[float]) -> Tensor:
+    if sigma is None:
+        sigma = [ksize * 0.15 + 0.35 for ksize in kernel_size]
+
+    if sigma is not None and not isinstance(sigma, (int, float, list, tuple)):
+        raise TypeError(
+            f"sigma should be either float or sequence of floats. Got {type(sigma)}"
+        )
+    if isinstance(sigma, (int, float)):
+        sigma = [float(sigma), float(sigma)]
+    if isinstance(sigma, (list, tuple)) and len(sigma) == 1:
+        sigma = [sigma[0], sigma[0]]
+    if len(sigma) != 2:
+        raise ValueError(
+            f"If sigma is a sequence, its length should be 2. Got {len(sigma)}"
+        )
+    for s in sigma:
+        if s <= 0.0:
+            raise ValueError(f"sigma should have positive values. Got {sigma}")
+    # print(f"img: {img}")
+    # if not (isinstance(img, Tensor)):
+    #     raise TypeError(f"img should be Tensor. Got {type(img)}")
+
+    dtype = img.dtype if paddle.is_floating_point(img) else paddle.float32
+    kernel = _get_gaussian_kernel2d(kernel_size, sigma, dtype=dtype)
+    kernel = kernel.expand([img.shape[-3], 1, kernel.shape[0], kernel.shape[1]])
+
+    img, need_cast, need_squeeze, out_dtype = _cast_squeeze_in(img, [kernel.dtype])
+
+    # padding = (left, right, top, bottom)
+    padding = [
+        kernel_size[0] // 2,
+        kernel_size[0] // 2,
+        kernel_size[1] // 2,
+        kernel_size[1] // 2,
+    ]
+    img = paddle.nn.functional.pad(img, padding, mode="reflect")
+    img = paddle.nn.functional.conv2d(img, kernel, groups=img.shape[-3])
+
+    img = _cast_squeeze_out(img, need_cast, need_squeeze, out_dtype)
+    return img

examples/data_efficient_nopt/pretrain_basic.py

@@ -23,6 +23,7 @@
 # from dadaptation import DAdaptAdan
 from einops import rearrange
 from models.fno import fno_pretrain
+from models.gaussian_blur import gaussian_blur
 from models.vmae import build_vmae
 
 # from paddle import DataParallel
@@ -411,11 +412,11 @@ def train_one_epoch(self):
                                 int((sigma * 4 + 1) / 2) * 2 + 1,
                                 (_inp.shape[2] // 2) * 2 - 1,
                             )
-                            # if _kernel >= 2:
-                            #     # [https://github.com/PaddlePaddle/Paddle/issues/26568]
-                            #     _inp = functional.gaussian_blur(
-                            #         _inp, kernel_size=[_kernel, _kernel], sigma=sigma
-                            #     )
+                            if _kernel >= 2:
+                                # [https://github.com/PaddlePaddle/Paddle/issues/26568]
+                                _inp = gaussian_blur(
+                                    _inp, kernel_size=[_kernel, _kernel], sigma=sigma
+                                )
                             inp_blur.append(_inp)
                         inp_blur = paddle.stack(inp_blur, axis=0)
                     else: