[Fix]fix data transform error (#779)

* [Fix]fix data transform error

* [Update]add 'auto_collation' for some examples

Author: lijialin03

docs/zh/examples/tempoGAN.md

@@ -154,9 +154,9 @@ examples/tempoGAN/tempoGAN.py:78:94
 
 本问题采用无监督学习的方式，虽然不是以监督学习方式进行训练，但此处仍然可以采用监督约束 `SupervisedConstraint`，在定义约束之前，需要给监督约束指定文件路径等数据读取配置，因为 tempoGAN 属于自监督学习，数据集中没有标签数据，而是使用一部分输入数据作为 `label`，因此需要设置约束的 `output_expr`。
 
-``` py linenums="122"
+``` py linenums="123"
 --8<--
-examples/tempoGAN/tempoGAN.py:122:125
+examples/tempoGAN/tempoGAN.py:123:126
 --8<--
 ```
 
@@ -166,7 +166,7 @@ examples/tempoGAN/tempoGAN.py:122:125
 
 ``` py linenums="98"
 --8<--
-examples/tempoGAN/tempoGAN.py:98:127
+examples/tempoGAN/tempoGAN.py:98:129
 --8<--
 ```
 
@@ -177,7 +177,9 @@ examples/tempoGAN/tempoGAN.py:98:127
 3. `label`： Array 类型的标签数据；
 4. `transforms`： 所有数据 transform 方法，此处 `FunctionalTransform` 为PaddleScience 预留的自定义数据 transform 类，该类支持编写代码时自定义输入数据的 transform，具体代码请参考 [自定义 loss 和 data transform](#38)；
 
-`batch_size` 字段表示 batch的大小；
+`auto_collation` 字段表示允许 BatchSampler 自动排序；
+
+`batch_size` 字段表示 batch 的大小；
 
 `sampler` 字段表示采样方法，其中各个字段表示：
 
@@ -193,27 +195,27 @@ examples/tempoGAN/tempoGAN.py:98:127
 
 在约束构建完毕之后，以我们刚才的命名为关键字，封装到一个字典中，方便后续访问，由于本问题设置了`use_spatialdisc` 和 `use_tempodisc`，导致 Generator 的部分约束不一定存在，因此先封装一定存在的约束到字典中，当其余约束存在时，在向字典中添加约束元素。
 
-``` py linenums="129"
+``` py linenums="130"
 --8<--
-examples/tempoGAN/tempoGAN.py:129:160
+examples/tempoGAN/tempoGAN.py:130:162
 --8<--
 ```
 
 #### 3.6.2 Discriminator 的约束
 
-``` py linenums="164"
+``` py linenums="166"
 --8<--
-examples/tempoGAN/tempoGAN.py:164:201
+examples/tempoGAN/tempoGAN.py:166:204
 --8<--
 ```
 
 各个参数含义与[Generator 的约束](#361)相同。
 
 #### 3.6.3 Discriminator_tempo 的约束
 
-``` py linenums="205"
+``` py linenums="208"
 --8<--
-examples/tempoGAN/tempoGAN.py:205:244
+examples/tempoGAN/tempoGAN.py:208:248
 --8<--
 ```
 
@@ -279,9 +281,9 @@ examples/tempoGAN/functions.py:430:488
 
 完成上述设置之后，首先需要将上述实例化的对象按顺序传递给 `ppsci.solver.Solver`，然后启动训练。
 
-``` py linenums="247"
+``` py linenums="251"
 --8<--
-examples/tempoGAN/tempoGAN.py:247:258
+examples/tempoGAN/tempoGAN.py:251:262
 --8<--
 ```
 
@@ -293,15 +295,15 @@ examples/tempoGAN/tempoGAN.py:247:258
 
 训练中仅在特定 `Epoch` 保存特定图片的目标结果和模型输出结果，训练结束后针对最后一个 `Epoch` 的输出结果进行一次评估，以便直观评价模型优化效果。不使用 PaddleScience 中内置的评估器，也不在训练过程中进行评估:
 
-``` py linenums="287"
+``` py linenums="291"
 --8<--
-examples/tempoGAN/tempoGAN.py:287:293
+examples/tempoGAN/tempoGAN.py:291:297
 --8<--
 ```
 
-``` py linenums="307"
+``` py linenums="311"
 --8<--
-examples/tempoGAN/tempoGAN.py:307:323
+examples/tempoGAN/tempoGAN.py:311:327
 --8<--
 ```
 
@@ -311,17 +313,17 @@ examples/tempoGAN/tempoGAN.py:307:323
 
 本问题的评估指标为，将模型输出的超分结果与实际高分辨率图片做对比，使用三个指标 MSE(Mean-Square Error) 、PSNR(Peak Signal-to-Noise Ratio) 、SSIM(Structural SIMilarity) 来评价图片相似度。因此没有使用 PaddleScience 中的内置评估器，也没有 `Solver.eval()` 过程。
 
-``` py linenums="326"
+``` py linenums="330"
 --8<--
-examples/tempoGAN/tempoGAN.py:326:406
+examples/tempoGAN/tempoGAN.py:330:410
 --8<--
 ```
 
 另外，其中：
 
-``` py linenums="396"
+``` py linenums="400"
 --8<--
-examples/tempoGAN/tempoGAN.py:396:403
+examples/tempoGAN/tempoGAN.py:400:407
 --8<--
 ```
 

docs/zh/examples/topopt.md

@@ -115,7 +115,7 @@ examples/topopt/topopt.py:36:38
 
 ``` py linenums="102"
 --8<--
-examples/topopt/functions.py:102:133
+examples/topopt/functions.py:102:135
 --8<--
 ```
 
@@ -125,7 +125,7 @@ examples/topopt/functions.py:102:133
 
 ``` py linenums="50"
 --8<--
-examples/topopt/topopt.py:50:75
+examples/topopt/topopt.py:50:76
 --8<--
 ```
 
@@ -136,7 +136,7 @@ examples/topopt/topopt.py:50:75
 3. `label`： 标签变量字典：`{"label_name": label_dataset}`；
 4. `transforms`： 数据集预处理配，其中 `"FunctionalTransform"` 为用户自定义的预处理方式。
 
-读取配置中 `"batch_size"` 字段表示训练时指定的批大小，`"sampler"` 字段表示 dataloader 的相关采样配置。
+读取配置中 `auto_collation` 字段表示允许 BatchSampler 自动排序， `batch_size` 字段表示训练时指定的批大小，`sampler` 字段表示 dataloader 的相关采样配置。
 
 第二个参数是损失函数，这里使用[自定义损失](#381)，通过 `cfg.vol_coeff` 确定损失公式中 $\beta$ 对应的值。
 
@@ -194,9 +194,9 @@ $$
 
 loss 构建代码如下：
 
-``` py linenums="263"
+``` py linenums="264"
 --8<--
-examples/topopt/topopt.py:263:274
+examples/topopt/topopt.py:264:275
 --8<--
 ```
 
@@ -215,9 +215,9 @@ $$
 其中 $n_{0} = w_{00} + w_{01}$ ， $n_{1} = w_{10} + w_{11}$ ，$w_{tp}$ 表示实际是 $t$ 类且被预测为 $p$ 类的像素点的数量
 metric 构建代码如下：
 
-``` py linenums="277"
+``` py linenums="278"
 --8<--
-examples/topopt/topopt.py:277:317
+examples/topopt/topopt.py:278:318
 --8<--
 ```
 
@@ -233,9 +233,9 @@ examples/topopt/conf/topopt.yaml:29:31
 
 训练代码如下：
 
-``` py linenums="77"
+``` py linenums="78"
 --8<--
-examples/topopt/topopt.py:77:111
+examples/topopt/topopt.py:78:111
 --8<--
 ```
 
@@ -249,7 +249,7 @@ examples/topopt/topopt.py:77:111
 
 ``` py linenums="218"
 --8<--
-examples/topopt/topopt.py:218:245
+examples/topopt/topopt.py:218:246
 --8<--
 ```
 

examples/tempoGAN/tempoGAN.py

@@ -111,6 +111,7 @@ def train(cfg: DictConfig):
                     },
                 ),
             },
+            "auto_collation": True,
             "batch_size": cfg.TRAIN.batch_size.sup_constraint,
             "sampler": {
                 "name": "BatchSampler",
@@ -143,6 +144,7 @@ def train(cfg: DictConfig):
                         },
                     ),
                 },
+                "auto_collation": True,
                 "batch_size": int(cfg.TRAIN.batch_size.sup_constraint // 3),
                 "sampler": {
                     "name": "BatchSampler",
@@ -188,6 +190,7 @@ def train(cfg: DictConfig):
                         },
                     ),
                 },
+                "auto_collation": True,
                 "batch_size": cfg.TRAIN.batch_size.sup_constraint,
                 "sampler": {
                     "name": "BatchSampler",
@@ -229,6 +232,7 @@ def train(cfg: DictConfig):
                         },
                     ),
                 },
+                "auto_collation": True,
                 "batch_size": int(cfg.TRAIN.batch_size.sup_constraint // 3),
                 "sampler": {
                     "name": "BatchSampler",

examples/topopt/functions.py

@@ -113,6 +113,8 @@ def augmentation(
     """
     inputs = input_dict["input"]
     labels = label_dict["output"]
+    assert len(inputs.shape) == 3
+    assert len(labels.shape) == 3
 
     # random horizontal flip
     if np.random.random() > 0.5:
@@ -125,7 +127,7 @@ def augmentation(
     # random 90* rotation
     if np.random.random() > 0.5:
         new_perm = list(range(len(inputs.shape)))
-        new_perm[1], new_perm[2] = new_perm[2], new_perm[1]
+        new_perm[-2], new_perm[-1] = new_perm[-1], new_perm[-2]
         inputs = np.transpose(inputs, new_perm)
         labels = np.transpose(labels, new_perm)
 

examples/topopt/topopt.py

@@ -62,6 +62,7 @@ def train(cfg: DictConfig):
                     },
                 ),
             },
+            "auto_collation": True,
             "batch_size": cfg.TRAIN.batch_size,
             "sampler": {
                 "name": "BatchSampler",
@@ -76,7 +77,6 @@ def train(cfg: DictConfig):
 
     # train models for 4 cases
     for sampler_key, num in cfg.CASE_PARAM:
-
         # initialize SIMP iteration stop time sampler
         SIMP_stop_point_sampler = func_module.generate_sampler(sampler_key, num)
 
@@ -229,6 +229,7 @@ def evaluate_model(
                         },
                     ),
                 },
+                "auto_collation": True,
                 "batch_size": cfg.EVAL.batch_size,
                 "sampler": {
                     "name": "BatchSampler",

examples/topopt/topoptmodel.py

@@ -39,7 +39,7 @@ class TopOptNN(ppsci.arch.UNetEx):
 
     Examples:
         >>> import ppsci
-        >>> model = ppsci.arch.ppsci.arch.UNetEx("input", "output", 2, 1, 3, (16, 32, 64), 2, lambda: 1, Flase, False)
+        >>> model = ppsci.arch.ppsci.arch.TopOptNN("input", "output", 2, 1, 3, (16, 32, 64), 2, lambda: 1, Flase, False)
     """
 
     def __init__(