fix uniform range for dit stage (#1215)

* fix uniform range for dit stage

* fix for ddp logging block

Author: HydrogenSulfate

docs/zh/examples/fundiff.md

@@ -34,7 +34,7 @@
 
 | 预训练模型  | 指标 |
 |:--| :--|
-| [fundiff_turbulence_mass_transfer_dit_pretrained.pdparams](https://paddle-org.bj.bcebos.com/paddlescience/models/fundiff/fundiff_turbulence_mass_transfer_dit_pretrained.pdparams) | Mean relative p error: 0.0651<br>Max relative p error: 0.1329<br>Min relative p error: 0.0345<br>Std relative p error: 0.0300<br>Mean relative sdf error: 0.0684<br>Max relative sdf error: 0.1443<br>Min relative sdf error: 0.0335<br>Std relative sdf error: 0.0377 |
+| [fundiff_turbulence_mass_transfer_dit_pretrained.pdparams](https://paddle-org.bj.bcebos.com/paddlescience/models/fundiff/fundiff_turbulence_mass_transfer_dit_pretrained.pdparams) | Mean relative p error: 0.066<br>Max relative p error: 0.159<br>Min relative p error: 0.029<br>Std relative p error: 0.027<br>Mean relative sdf error: 0.085<br>Max relative sdf error: 0.307<br>Min relative sdf error: 0.022<br>Std relative sdf error: 0.0499 |
 
 ## 1. 背景简介
 

examples/fundiff/conf/diffusion.yaml

@@ -61,7 +61,7 @@ FAE:
     num_mlp_layers: 2
     out_dim: 1
     layer_norm_eps: 1e-05
-  pretrained_model_path: null
+  pretrained_model_path: ???
 
 DIT:
   # input_keys: [u, v, p, sdf]

examples/fundiff/main.py

@@ -66,7 +66,11 @@ def forward(self, batch: Dict[str, paddle.Tensor]):
                 z_sdf = self.encoder(sdf)
                 z_1 = paddle.concat([z_p, z_sdf], axis=-1)
                 z_0 = paddle.randn(z_1.shape)  # (b, 200, 512)
-                t = paddle.uniform([z_1.shape[0], *[1 for _ in range(z_1.ndim - 1)]])
+                t = paddle.uniform(
+                    [z_1.shape[0], *[1 for _ in range(z_1.ndim - 1)]],
+                    min=0.0,
+                    max=1.0,
+                )
                 z_t = t * (z_1 - z_0) + z_0
                 v_t = z_1 - z_0
             else:

ppsci/utils/logger.py

@@ -213,22 +213,19 @@ def scalar(
         wandb_writer (Optional[wandb.run]): Run object of WandB to record metrics. Defaults to None.
         tbd_writer (Optional[tbd.SummaryWriter]): Run object of WandB to record metrics. Defaults to None.
     """
-    if vdl_writer is not None:
-        with misc.RankZeroOnly() as is_master:
-            if is_master:
-                for name, value in metric_dict.items():
-                    vdl_writer.add_scalar(name, value, step)
-
-    if wandb_writer is not None:
-        with misc.RankZeroOnly() as is_master:
-            if is_master:
-                wandb_writer.log({"step": step, **metric_dict})
-
-    if tbd_writer is not None:
-        with misc.RankZeroOnly() as is_master:
-            if is_master:
-                for name, value in metric_dict.items():
-                    tbd_writer.add_scalar(name, value, global_step=step)
+    with misc.RankZeroOnly() as is_master:
+        if vdl_writer is not None and is_master:
+            for name, value in metric_dict.items():
+                vdl_writer.add_scalar(name, value, step)
+
+    with misc.RankZeroOnly() as is_master:
+        if wandb_writer is not None and is_master:
+            wandb_writer.log({"step": step, **metric_dict})
+
+    with misc.RankZeroOnly() as is_master:
+        if tbd_writer is not None and is_master:
+            for name, value in metric_dict.items():
+                tbd_writer.add_scalar(name, value, global_step=step)
 
 
 def advertise():