[TIPC] add tipc benchmark for msvsr (#672)

Birdylx · web-flow · commit eebf94d99507 · 2022-08-19T14:41:57.000+08:00
* add tipc benchmark for msvsr

* update tipc readme img
diff --git a/ppgan/engine/trainer.py b/ppgan/engine/trainer.py
@@ -32,6 +32,7 @@
 
 
 class IterLoader:
+
     def __init__(self, dataloader):
         self._dataloader = dataloader
         self.iter_loader = iter(self._dataloader)
@@ -79,6 +80,7 @@ class Trainer:
     #                     |                                    ||
     #         save checkpoint (model.nets)                     \/
     """
+
     def __init__(self, cfg):
         # base config
         self.logger = logging.getLogger(__name__)
@@ -181,6 +183,22 @@ def train(self):
 
         iter_loader = IterLoader(self.train_dataloader)
 
+        # use amp
+        if self.cfg.amp:
+            self.logger.info('use AMP to train. AMP level = {}'.format(
+                self.cfg.amp_level))
+            assert self.cfg.model.name == 'MultiStageVSRModel', "AMP only support msvsr model"
+            scaler = paddle.amp.GradScaler(init_loss_scaling=1024)
+            # need to decorate model and optim if amp_level == 'O2'
+            if self.cfg.amp_level == 'O2':
+                # msvsr has only one generator and one optimizer
+                self.model.nets['generator'], self.optimizers[
+                    'optim'] = paddle.amp.decorate(
+                        models=self.model.nets['generator'],
+                        optimizers=self.optimizers['optim'],
+                        level='O2',
+                        save_dtype='float32')
+
         # set model.is_train = True
         self.model.setup_train_mode(is_train=True)
         while self.current_iter < (self.total_iters + 1):
@@ -195,7 +213,12 @@ def train(self):
             # unpack data from dataset and apply preprocessing
             # data input should be dict
             self.model.setup_input(data)
-            self.model.train_iter(self.optimizers)
+
+            if self.cfg.amp:
+                self.model.train_iter_amp(self.optimizers, scaler,
+                                          self.cfg.amp_level)  # amp train
+            else:
+                self.model.train_iter(self.optimizers)  # norm train
 
             batch_cost_averager.record(
                 time.time() - step_start_time,
diff --git a/ppgan/models/msvsr_model.py b/ppgan/models/msvsr_model.py
@@ -30,6 +30,7 @@ class MultiStageVSRModel(BaseSRModel):
     Paper:
         PP-MSVSR: Multi-Stage Video Super-Resolution, 2021
     """
+
     def __init__(self, generator, fix_iter, pixel_criterion=None):
         """Initialize the PP-MSVSR class.
 
@@ -96,6 +97,48 @@ def train_iter(self, optims=None):
 
         self.current_iter += 1
 
+    # amp train with brute force implementation, maybe decorator can simplify this
+    def train_iter_amp(self, optims=None, scaler=None, amp_level='O1'):
+        optims['optim'].clear_grad()
+        if self.fix_iter:
+            if self.current_iter == 1:
+                print('Train MSVSR with fixed spynet for', self.fix_iter,
+                      'iters.')
+                for name, param in self.nets['generator'].named_parameters():
+                    if 'spynet' in name:
+                        param.trainable = False
+            elif self.current_iter >= self.fix_iter + 1 and self.flag:
+                print('Train all the parameters.')
+                for name, param in self.nets['generator'].named_parameters():
+                    param.trainable = True
+                    if 'spynet' in name:
+                        param.optimize_attr['learning_rate'] = 0.25
+                self.flag = False
+                for net in self.nets.values():
+                    net.find_unused_parameters = False
+
+        # put loss computation in amp context
+        with paddle.amp.auto_cast(enable=True, level=amp_level):
+            output = self.nets['generator'](self.lq)
+            if isinstance(output, (list, tuple)):
+                out_stage2, output = output
+                loss_pix_stage2 = self.pixel_criterion(out_stage2, self.gt)
+                self.losses['loss_pix_stage2'] = loss_pix_stage2
+            self.visual_items['output'] = output[:, 0, :, :, :]
+            # pixel loss
+            loss_pix = self.pixel_criterion(output, self.gt)
+            self.losses['loss_pix'] = loss_pix
+
+            self.loss = sum(_value for _key, _value in self.losses.items()
+                            if 'loss_pix' in _key)
+        scaled_loss = scaler.scale(self.loss)
+        self.losses['loss'] = scaled_loss
+
+        scaled_loss.backward()
+        scaler.minimize(optims['optim'], scaled_loss)
+
+        self.current_iter += 1
+
     def test_iter(self, metrics=None):
         self.gt = self.gt.cpu()
         self.nets['generator'].eval()
diff --git a/ppgan/utils/options.py b/ppgan/utils/options.py
@@ -45,9 +45,9 @@ def parse_args():
                         default=False,
                         help='skip validation during training')
     # config options
-    parser.add_argument("-o", 
-                        "--opt", 
-                        nargs='+', 
+    parser.add_argument("-o",
+                        "--opt",
+                        nargs='+',
                         help="set configuration options")
 
     #for inference
@@ -60,19 +60,31 @@ def parse_args():
                         help="path to reference images")
     parser.add_argument("--model_path", default=None, help="model for loading")
 
-    # for profiler                
-    parser.add_argument('-p',
-                        '--profiler_options',
-                        type=str,
-                        default=None,
-                        help='The option of profiler, which should be in format \"key1=value1;key2=value2;key3=value3\".'
+    # for profiler
+    parser.add_argument(
+        '-p',
+        '--profiler_options',
+        type=str,
+        default=None,
+        help=
+        'The option of profiler, which should be in format \"key1=value1;key2=value2;key3=value3\".'
     )
     # fix random numbers by setting seed
     parser.add_argument('--seed',
                         type=int,
                         default=None,
-                        help='fix random numbers by setting seed\".'
-    )
+                        help='fix random numbers by setting seed\".')
+
+    # add for amp training
+    parser.add_argument('--amp',
+                        action='store_true',
+                        default=False,
+                        help='whether to enable amp training')
+    parser.add_argument('--amp_level',
+                        type=str,
+                        default='O1',
+                        choices=['O1', 'O2'],
+                        help='level of amp training; O2 represent pure fp16')
     args = parser.parse_args()
 
     return args
diff --git a/ppgan/utils/setup.py b/ppgan/utils/setup.py
@@ -19,6 +19,7 @@
 import random
 from .logger import setup_logger
 
+
 def setup(args, cfg):
     if args.evaluate_only:
         cfg.is_train = False
@@ -44,10 +45,13 @@ def setup(args, cfg):
         paddle.set_device('gpu')
     else:
         paddle.set_device('cpu')
-    
+
     if args.seed:
         paddle.seed(args.seed)
         random.seed(args.seed)
-        np.random.seed(args.seed)    
+        np.random.seed(args.seed)
         paddle.framework.random._manual_program_seed(args.seed)
-    
+
+    # add amp and amp_level args into cfg
+    cfg['amp'] = args.amp
+    cfg['amp_level'] = args.amp_level
diff --git a/test_tipc/README.md b/test_tipc/README.md
@@ -57,9 +57,8 @@ test_tipc/
 
 ### 测试流程
 使用本工具，可以测试不同功能的支持情况，以及预测结果是否对齐，测试流程如下：
-<div align="center">
-    <img src="docs/test.png" width="800">
-</div>
+
+![img](https://user-images.githubusercontent.com/79366697/185377097-a0f852a8-2d78-45ae-84ba-ae71b799d738.png)
 
 1. 运行prepare.sh准备测试所需数据和模型；
 2. 运行要测试的功能对应的测试脚本`test_*.sh`，产出log，由log可以看到不同配置是否运行成功；
@@ -72,4 +71,4 @@ test_tipc/
 <a name="more"></a>
 #### 更多教程
 各功能测试中涉及混合精度、裁剪、量化等训练相关，及mkldnn、Tensorrt等多种预测相关参数配置，请点击下方相应链接了解更多细节和使用教程：  
-[test_train_inference_python 使用](docs/test_train_inference_python.md)  
+- [test_train_inference_python 使用](docs/test_train_inference_python.md): 测试基于Python的模型训练、评估、推理等基本功能
diff --git a/test_tipc/benchmark_train.sh b/test_tipc/benchmark_train.sh
@@ -4,15 +4,15 @@ source test_tipc/common_func.sh
 # set env
 python=python
 export model_branch=`git symbolic-ref HEAD 2>/dev/null | cut -d"/" -f 3`
-export model_commit=$(git log|head -n1|awk '{print $2}') 
+export model_commit=$(git log|head -n1|awk '{print $2}')
 export str_tmp=$(echo `pip list|grep paddlepaddle-gpu|awk -F ' ' '{print $2}'`)
 export frame_version=${str_tmp%%.post*}
 export frame_commit=$(echo `${python} -c "import paddle;print(paddle.version.commit)"`)
 
-# run benchmark sh 
+# run benchmark sh
 # Usage:
 # bash run_benchmark_train.sh config.txt params
-# or 
+# or
 # bash run_benchmark_train.sh config.txt
 
 function func_parser_params(){
@@ -100,6 +100,7 @@ for _flag in ${flags_list[*]}; do
 done
 
 # set log_name
+BENCHMARK_ROOT=./ # self-test only
 repo_name=$(get_repo_name )
 SAVE_LOG=${BENCHMARK_LOG_DIR:-$(pwd)}   # */benchmark_log
 mkdir -p "${SAVE_LOG}/benchmark_log/"
@@ -149,11 +150,11 @@ else
 fi
 
 IFS="|"
-for batch_size in ${batch_size_list[*]}; do 
+for batch_size in ${batch_size_list[*]}; do
     for precision in ${fp_items_list[*]}; do
         for device_num in ${device_num_list[*]}; do
             # sed batchsize and precision
-            #func_sed_params "$FILENAME" "${line_precision}" "$precision"
+            func_sed_params "$FILENAME" "${line_precision}" "$precision"
             func_sed_params "$FILENAME" "${line_batchsize}" "$MODE=$batch_size"
             func_sed_params "$FILENAME" "${line_epoch}" "$MODE=$epoch"
             gpu_id=$(set_gpu_id $device_num)
@@ -162,7 +163,7 @@ for batch_size in ${batch_size_list[*]}; do
                 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"
-                func_sed_params "$FILENAME" "${line_gpuid}" "0"  # sed used gpu_id 
+                func_sed_params "$FILENAME" "${line_gpuid}" "0"  # sed used gpu_id
                 # set profile_option params
                 tmp=`sed -i "${line_profile}s/.*/${profile_option}/" "${FILENAME}"`
 
@@ -214,7 +215,7 @@ for batch_size in ${batch_size_list[*]}; do
                 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"
-                func_sed_params "$FILENAME" "${line_gpuid}" "$gpu_id"  # sed used gpu_id 
+                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 "
                 echo $cmd
@@ -244,4 +245,4 @@ for batch_size in ${batch_size_list[*]}; do
             fi
         done
     done
-done 
+done
diff --git a/test_tipc/configs/msvsr/train_amp_infer_python.txt b/test_tipc/configs/msvsr/train_amp_infer_python.txt
@@ -0,0 +1,53 @@
+===========================train_params===========================
+model_name:msvsr
+python:python3.7
+gpu_list:0
+##
+auto_cast:null
+total_iters:lite_train_lite_infer=10|lite_train_whole_infer=10|whole_train_whole_infer=200
+output_dir:./output/
+dataset.train.batch_size:lite_train_lite_infer=1|whole_train_whole_infer=1
+pretrained_model:null
+train_model_name:msvsr_reds*/*checkpoint.pdparams
+train_infer_img_dir:./data/msvsr_reds/test
+null:null
+##
+trainer:amp_train
+amp_train:tools/main.py --amp --amp_level O1 -c configs/msvsr_reds.yaml --seed 123 -o dataset.train.num_workers=0 log_config.interval=1 snapshot_config.interval=5 dataset.train.dataset.num_frames=2
+pact_train:null
+fpgm_train:null
+distill_train:null
+null:null
+null:null
+##
+===========================eval_params===========================
+eval:null
+null:null
+##
+===========================infer_params===========================
+--output_dir:./output/
+load:null
+norm_export:tools/export_model.py -c configs/msvsr_reds.yaml  --inputs_size="1,2,3,180,320"  --model_name inference --load
+quant_export:null
+fpgm_export:null
+distill_export:null
+export1:null
+export2:null
+inference_dir:inference
+train_model:./inference/msvsr/multistagevsrmodel_generator
+infer_export:null
+infer_quant:False
+inference:tools/inference.py --model_type msvsr -c configs/msvsr_reds.yaml --seed 123 -o dataset.test.num_frames=2 --output_path test_tipc/output/
+--device:cpu
+null:null
+null:null
+null:null
+null:null
+null:null
+--model_path:
+null:null
+null:null
+--benchmark:True
+null:null
+===========================infer_benchmark_params==========================
+random_infer_input:[{float32,[2,3,180,320]}]
diff --git a/test_tipc/configs/msvsr/train_infer_python.txt b/test_tipc/configs/msvsr/train_infer_python.txt
@@ -13,22 +13,22 @@ train_infer_img_dir:./data/msvsr_reds/test
 null:null
 ##
 trainer:norm_train
-norm_train:tools/main.py -c configs/msvsr_reds.yaml --seed 123 -o dataset.train.num_workers=0 log_config.interval=1 snapshot_config.interval=5 dataset.train.dataset.num_frames=2 
+norm_train:tools/main.py -c configs/msvsr_reds.yaml --seed 123 -o log_config.interval=2 snapshot_config.interval=50 dataset.train.dataset.num_frames=15
 pact_train:null
 fpgm_train:null
 distill_train:null
 null:null
 null:null
 ##
-===========================eval_params=========================== 
+===========================eval_params===========================
 eval:null
 null:null
 ##
 ===========================infer_params===========================
 --output_dir:./output/
 load:null
-norm_export:tools/export_model.py -c configs/msvsr_reds.yaml  --inputs_size="1,2,3,180,320"  --model_name inference --load 
-quant_export:null 
+norm_export:tools/export_model.py -c configs/msvsr_reds.yaml  --inputs_size="1,2,3,180,320"  --model_name inference --load
+quant_export:null
 fpgm_export:null
 distill_export:null
 export1:null
@@ -49,5 +49,11 @@ null:null
 null:null
 --benchmark:True
 null:null
+===========================train_benchmark_params==========================
+batch_size:4
+fp_items:fp32
+total_iters:60
+--profiler_options:batch_range=[10,20];state=GPU;tracer_option=Default;profile_path=model.profile
+flags:null
 ===========================infer_benchmark_params==========================
 random_infer_input:[{float32,[2,3,180,320]}]
diff --git a/test_tipc/docs/benchmark_train.md b/test_tipc/docs/benchmark_train.md
diff --git a/test_tipc/docs/test.png b/test_tipc/docs/test.png
diff --git a/test_tipc/prepare.sh b/test_tipc/prepare.sh
diff --git a/test_tipc/test_train_inference_python.sh b/test_tipc/test_train_inference_python.sh
