[TIPC-Benchmark]Support @to_static traing for Benchmark for cyclegan (#739)

configs/cyclegan_horse2zebra.yaml

@@ -26,6 +26,8 @@ model:
   gan_criterion:
     name: GANLoss
     gan_mode: lsgan
+  # training model under @to_static
+  to_static: False
 
 export_model:
   - {name: 'netG_A', inputs_num: 1}

ppgan/models/__init__.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .base_model import BaseModel
+from .base_model import BaseModel, apply_to_static
 from .gan_model import GANModel
 from .cycle_gan_model import CycleGANModel
 from .pix2pix_model import Pix2PixModel

ppgan/models/base_model.py

@@ -19,9 +19,13 @@ import numpy as np
 from collections import OrderedDict
 from abc import ABC, abstractmethod
 
+from paddle.jit import to_static
+from paddle.static import InputSpec
+
 from .criterions.builder import build_criterion
 from ..solver import build_lr_scheduler, build_optimizer
 from ..utils.visual import tensor2img
+from ..utils.logger import get_logger
 
 
 class BaseModel(ABC):
@@ -217,3 +221,13 @@ class BaseModel(ABC):
                                                                 model_name),
                     model_filename="{}.pdmodel".format(model_name),
                     params_filename="{}.pdiparams".format(model_name))
+
+def apply_to_static(support_to_static, image_shape, model):
+    if support_to_static:
+        specs = None
+        if image_shape is not None:
+            specs = [InputSpec([None] + image_shape)]
+        model = to_static(model, input_spec=specs)
+        logger = get_logger('ppgan')
+        logger.info("Successfully to apply @to_static with specs: {}".format(specs))
+    return model
\ No newline at end of file

ppgan/models/cycle_gan_model.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import paddle
-from .base_model import BaseModel
+from .base_model import BaseModel, apply_to_static
 
 from .builder import MODELS
 from .generators.builder import build_generator
@@ -40,7 +40,9 @@ class CycleGANModel(BaseModel):
                  pool_size=50,
                  direction='a2b',
                  lambda_a=10.,
-                 lambda_b=10.):
+                 lambda_b=10.,
+                 to_static=False,
+                 image_shape=None):
         """Initialize the CycleGAN class.
 
         Args:
@@ -59,6 +61,9 @@ class CycleGANModel(BaseModel):
         # Code (vs. paper): G_A (G), G_B (F), D_A (D_Y), D_B (D_X)
         self.nets['netG_A'] = build_generator(generator)
         self.nets['netG_B'] = build_generator(generator)
+        # set @to_static for benchmark, skip this by default.
+        apply_to_static(to_static, image_shape, self.nets['netG_A'])
+        apply_to_static(to_static, image_shape, self.nets['netG_B'])
         init_weights(self.nets['netG_A'])
         init_weights(self.nets['netG_B'])
 
@@ -66,6 +71,9 @@ class CycleGANModel(BaseModel):
         if discriminator:
             self.nets['netD_A'] = build_discriminator(discriminator)
             self.nets['netD_B'] = build_discriminator(discriminator)
+            # set @to_static for benchmark, skip this by default.
+            apply_to_static(to_static, image_shape, self.nets['netD_A'])
+            apply_to_static(to_static, image_shape, self.nets['netD_B'])
             init_weights(self.nets['netD_A'])
             init_weights(self.nets['netD_B'])
 

test_tipc/benchmark_train.sh

@@ -67,6 +67,19 @@ FILENAME=$new_filename
 # MODE must be one of ['benchmark_train']
 MODE=$2
 PARAMS=$3
+REST_ARGS=$4
+
+to_static="d2sF"
+# parse "to_static" options and modify trainer into "to_static_trainer"
+if [ $REST_ARGS = "to_static" ] || [ $PARAMS = "to_static" ] ;then
+   to_static="d2sT"
+   sed -i 's/trainer:norm_train/trainer:to_static_train/g' $FILENAME
+   # clear PARAM contents
+   if [ $PARAMS = "to_static" ] ;then
+    PARAMS=""
+   fi
+fi
+
 IFS=$'\n'
 # parser params from train_benchmark.txt
 dataline=`cat $FILENAME`
@@ -161,7 +174,7 @@ for batch_size in ${batch_size_list[*]}; do
             if [ ${#gpu_id} -le 1 ];then
                 log_path="$SAVE_LOG/profiling_log"
                 mkdir -p $log_path
-                log_name="${repo_name}_${model_name}_bs${batch_size}_${precision}_${run_mode}_${device_num}_profiling"
+                log_name="${repo_name}_${model_name}_bs${batch_size}_${precision}_${run_mode}_${device_num}_${to_static}_profiling"
                 func_sed_params "$FILENAME" "${line_gpuid}" "0"  # sed used gpu_id
                 # set profile_option params
                 tmp=`sed -i "${line_profile}s/.*/${profile_option}/" "${FILENAME}"`
@@ -177,8 +190,8 @@ for batch_size in ${batch_size_list[*]}; do
                 speed_log_path="$SAVE_LOG/index"
                 mkdir -p $log_path
                 mkdir -p $speed_log_path
-                log_name="${repo_name}_${model_name}_bs${batch_size}_${precision}_${run_mode}_${device_num}_log"
-                speed_log_name="${repo_name}_${model_name}_bs${batch_size}_${precision}_${run_mode}_${device_num}_speed"
+                log_name="${repo_name}_${model_name}_bs${batch_size}_${precision}_${run_mode}_${device_num}_${to_static}_log"
+                speed_log_name="${repo_name}_${model_name}_bs${batch_size}_${precision}_${run_mode}_${device_num}_${to_static}_speed"
                 func_sed_params "$FILENAME" "${line_profile}" "null"  # sed profile_id as null
                 cmd="bash test_tipc/test_train_inference_python.sh ${FILENAME} benchmark_train > ${log_path}/${log_name} 2>&1 "
                 echo $cmd
@@ -212,8 +225,8 @@ for batch_size in ${batch_size_list[*]}; do
                 speed_log_path="$SAVE_LOG/index"
                 mkdir -p $log_path
                 mkdir -p $speed_log_path
-                log_name="${repo_name}_${model_name}_bs${batch_size}_${precision}_${run_mode}_${device_num}_log"
-                speed_log_name="${repo_name}_${model_name}_bs${batch_size}_${precision}_${run_mode}_${device_num}_speed"
+                log_name="${repo_name}_${model_name}_bs${batch_size}_${precision}_${run_mode}_${device_num}_${to_static}_log"
+                speed_log_name="${repo_name}_${model_name}_bs${batch_size}_${precision}_${run_mode}_${device_num}_${to_static}_speed"
                 func_sed_params "$FILENAME" "${line_gpuid}" "$gpu_id"  # sed used gpu_id
                 func_sed_params "$FILENAME" "${line_profile}" "null"  # sed --profile_option as null
                 cmd="bash test_tipc/test_train_inference_python.sh ${FILENAME} benchmark_train > ${log_path}/${log_name} 2>&1 "

test_tipc/configs/CycleGAN/train_infer_python.txt

@@ -17,7 +17,7 @@ norm_train:tools/main.py -c configs/cyclegan_horse2zebra.yaml --seed 123 -o log_
 pact_train:null
 fpgm_train:null
 distill_train:null
-null:null
+to_static_train:-o model.to_static=True
 null:null
 ##
 ===========================eval_params=========================== 

test_tipc/test_train_inference_python.sh

@@ -33,8 +33,8 @@ trainer_list=$(func_parser_value "${lines[14]}")
 trainer_norm=$(func_parser_key "${lines[15]}")
 norm_trainer=$(func_parser_value "${lines[15]}")
 
-trainer_key1=$(func_parser_key "${lines[19]}")
-trainer_value1=$(func_parser_value "${lines[19]}")
+to_static_key=$(func_parser_key "${lines[19]}")
+to_static_trainer=$(func_parser_value "${lines[19]}")
 trainer_key2=$(func_parser_key "${lines[20]}")
 trainer_value2=$(func_parser_value "${lines[20]}")
 
@@ -218,8 +218,16 @@ else
             fi
             for trainer in ${trainer_list[*]}; do
                 flag_quant=False
-                run_train=${norm_trainer}
-                run_export=${norm_export}
+                # In case of @to_static, we re-used norm_traier,
+                # but append "-o Global.to_static=True" for config
+                # to trigger "apply_to_static" logic in 'engine.py'
+                if [ ${trainer} = "${to_static_key}" ]; then
+                    run_train="${norm_trainer}  ${to_static_trainer}"
+                    run_export=${norm_export}
+                else
+                    run_train=${norm_trainer}
+                    run_export=${norm_export}
+                fi
 
                 if [ ${run_train} = "null" ]; then
                     continue