Refine some configurations in TSN model (#4853)

Author: LiuChaoXD

dygraph/tsn/README.md

@@ -1,5 +1,5 @@
 # TSN 视频分类模型
-本目录下为基于PaddlePaddle 动态图实现的TSN视频分类模型。模型支持PaddlePaddle Fluid 1.8, GPU, Linux。
+本目录下为基于PaddlePaddle 动态图实现的TSN视频分类模型。模型支持PaddlePaddle Fluid 2.0, GPU, Linux。
 
 ---
 ## 内容
@@ -50,7 +50,7 @@ Temporal Segment Network (TSN) 是视频分类领域经典的基于2D-CNN的解
 
 ## 数据准备
 
-TSN的训练数据采用UCF101动作识别数据集。数据下载及处理请参考[数据说明](./data/dataset/ucf101/README.md)。数据处理完成后，会在`./data/dataset/ucf101/`目录下，生成一下文件：
+TSN的训练数据采用UCF101动作识别数据集。数据下载及处理请参考[数据说明](./data/dataset/ucf101/README.md)。数据处理完成后，会在`./data/dataset/ucf101/`目录下，生成以下文件：
 - `videos/` ： 用于存放UCF101数据的视频文件。
 - `rawframes/` ： 用于存放UCF101视频文件的frame数据。
 - `annotations/` ： 用于存储UCF101数据集的标注文件。
@@ -60,16 +60,21 @@ TSN的训练数据采用UCF101动作识别数据集。数据下载及处理请
 
 
 ## 模型训练
+TSN模型训练，需要加载基于imagenet pretrain的ResNet50参数。可通过输入如下命令下载（默认权重文件会存放在当前目前下`./ResNet50_pretrained/`）：
+```bash
+bash download_pretrain.sh
+```
 
-TSN模型支持输入数据为video和frame格式。数据准备完毕后，可以通过如下方式启动不同格式的训练。
+TSN模型支持输入数据为video和frame格式。数据以及预训练参数准备完毕后，可以通过如下方式启动不同格式的训练。
 
 1. 多卡训练（输入为frame格式）
 ```bash
 bash multi_gpus_run.sh ./multi_tsn_frame.yaml
 ```
 多卡训练所使用的gpu可以通过如下方式设置：
 - 修改`multi_gpus_run.sh` 中 `export CUDA_VISIBLE_DEVICES=0,1,2,3`（默认为0,1,2,3表示使用0，1，2，3卡号的gpu进行训练）
-- 注意：多卡、frame格式的训练参数配置文件为`multi_tsn_frame.yaml`。若修改了batchsize则学习率也要做相应的修改，规则为大batchsize用大lr，即同倍数增长缩小关系。例如，默认四卡batchsize=128，lr=0.001，若batchsize=64，lr=0.0005。
+- 若需要修改预训练权重文件的加载路径，可在`multi_gpus_run.sh`中修改`pretrain`参数（默认`pretrain="./ResNet50_pretrained/"`）
+- 注意：多卡、frame格式的训练参数配置文件为`multi_tsn_frame.yaml`。若修改了batchsize则学习率也要做相应的修改，规则为大batchsize用大lr，即同倍数增大缩小关系。例如，默认四卡batchsize=128，lr=0.001，若batchsize=64，lr=0.0005。
 
 2. 多卡训练（输入为video格式）
 ```bash
@@ -85,7 +90,8 @@ bash single_gpu_run.sh ./single_tsn_frame.yaml
 ```
 单卡训练所使用的gpu可以通过如下方式设置：
 - 修改 `single_gpu_run.sh` 中的 `export CUDA_VISIBLE_DEVICES=0` （表示使用gpu 0 进行模型训练）
-- 注意：单卡、frame格式的训练参数配置文件为`single_tsn_frame.yaml`。若修改了batchsize则学习率也要做相应的修改，规则为大batchsize用大lr，即同倍数增长缩小关系。默认单卡batchsize=64，lr=0.0005；若batchsize=32，lr=0.00025。
+- 若需要修改预训练权重文件的加载路径，可在`single_gpu_run.sh`中修改`pretrain`参数（默认`pretrain="./ResNet50_pretrained/"`）
+- 注意：单卡、frame格式的训练参数配置文件为`single_tsn_frame.yaml`。若修改了batchsize则学习率也要做相应的修改，规则为大batchsize用大lr，即同倍数增长缩小关系。默认单卡batchsize=32，lr=0.00025；若batchsize=64，lr=0.0005。
 
 4. 单卡训练（输入为video格式）
 ```bash

dygraph/tsn/download_pretrain.sh

@@ -0,0 +1,4 @@
+#!/bin/bash
+wget https://paddlemodels.bj.bcebos.com/video_classification/ResNet50_pretrained.tar.gz
+tar -xzvf ResNet50_pretrained.tar.gz
+rm -rf ResNet50_pretrained.tar.gz

dygraph/tsn/multi_gpus_run.sh

@@ -1,11 +1,11 @@
 configs=$1
-pretrain="" # set pretrain model path if needed
+pretrain="./ResNet50_pretrained/" # set pretrain model path if needed
 resume="" # set checkpoints model path if u want to resume training
 save_dir=""
 use_gpu=True
 use_data_parallel=True
 
-export CUDA_VISIBLE_DEVICES=4,5,6,7
+export CUDA_VISIBLE_DEVICES=0,1,2,3
 
 
 echo $mode "TSN" $configs  $resume $pretrain

dygraph/tsn/multi_tsn_frame.yaml

@@ -13,8 +13,8 @@ TRAIN:
     epoch: 80
     short_size: 256
     target_size: 224
-    num_reader_threads: 12
-    buf_size: 1024
+    num_reader_threads: 16
+    buf_size: 256
     batch_size: 128
     use_gpu: True
     filelist: "./data/dataset/ucf101/ucf101_train_split_1_rawframes.txt"
@@ -28,15 +28,15 @@ TRAIN:
 VALID:
     short_size: 256
     target_size: 224
-    num_reader_threads: 12
-    buf_size: 1024
+    num_reader_threads: 16
+    buf_size: 256
     batch_size: 128
     filelist: "./data/dataset/ucf101/ucf101_val_split_1_rawframes.txt"
 
 TEST:
     short_size: 256
     target_size: 224
-    num_reader_threads: 12
-    buf_size: 1024
-    batch_size: 64
-    filelist: "./data/dataset/ucf101/ucf101_val_split_1_rawframes.txt"
+    num_reader_threads: 16
+    buf_size: 256
+    batch_size: 128
+    filelist: "./data/dataset/ucf101/ucf101_val_split_1_rawframes.txt"

dygraph/tsn/multi_tsn_video.yaml

@@ -13,8 +13,8 @@ TRAIN:
     epoch: 80
     short_size: 256
     target_size: 224
-    num_reader_threads: 12
-    buf_size: 1024
+    num_reader_threads: 16
+    buf_size: 256
     batch_size: 128
     use_gpu: True
     filelist: "./data/dataset/ucf101/ucf101_train_split_1_videos.txt"
@@ -28,15 +28,15 @@ TRAIN:
 VALID:
     short_size: 256
     target_size: 224
-    num_reader_threads: 12
-    buf_size: 1024
+    num_reader_threads: 16
+    buf_size: 256
     batch_size: 128
     filelist: "./data/dataset/ucf101/ucf101_val_split_1_videos.txt"
 
 TEST:
     short_size: 256
     target_size: 224
-    num_reader_threads: 12
-    buf_size: 1024
-    batch_size: 64
+    num_reader_threads: 16
+    buf_size: 256
+    batch_size: 128
     filelist: "./data/dataset/ucf101/ucf101_val_split_1_videos.txt"

dygraph/tsn/reader/__init__.py

@@ -1,5 +1,5 @@
 configs=$1
-pretrain="" # set pretrain model path if needed
+pretrain="./ResNet50_pretrained/" # set pretrain model path if needed
 resume="" # set checkpoints model path if u want to resume training
 save_dir=""
 use_gpu=True

dygraph/tsn/single_gpu_run.sh

@@ -13,12 +13,12 @@ TRAIN:
     epoch: 80
     short_size: 256
     target_size: 224
-    num_reader_threads: 12
-    buf_size: 1024
-    batch_size: 64
+    num_reader_threads: 8
+    buf_size: 64
+    batch_size: 32
     use_gpu: True
     filelist: "./data/dataset/ucf101/ucf101_train_split_1_rawframes.txt"
-    learning_rate: 0.0005
+    learning_rate: 0.00025
     learning_rate_decay: 0.1
     decay_epochs: [30, 60]
     l2_weight_decay: 1e-4
@@ -28,15 +28,15 @@ TRAIN:
 VALID:
     short_size: 256
     target_size: 224
-    num_reader_threads: 12
-    buf_size: 1024
-    batch_size: 128
+    num_reader_threads: 8
+    buf_size: 64
+    batch_size: 32
     filelist: "./data/dataset/ucf101/ucf101_val_split_1_rawframes.txt"
 
 TEST:
     short_size: 256
     target_size: 224
-    num_reader_threads: 12
-    buf_size: 1024
-    batch_size: 64
+    num_reader_threads: 8
+    buf_size: 64
+    batch_size: 32
     filelist: "./data/dataset/ucf101/ucf101_val_split_1_rawframes.txt"

dygraph/tsn/single_tsn_frame.yaml

@@ -13,12 +13,12 @@ TRAIN:
     epoch: 80
     short_size: 256
     target_size: 224
-    num_reader_threads: 12
-    buf_size: 1024
-    batch_size: 64
+    num_reader_threads: 8
+    buf_size: 64
+    batch_size: 32
     use_gpu: True
     filelist: "./data/dataset/ucf101/ucf101_train_split_1_videos.txt"
-    learning_rate: 0.0005
+    learning_rate: 0.00025
     learning_rate_decay: 0.1
     decay_epochs: [30, 60]
     l2_weight_decay: 1e-4
@@ -28,15 +28,15 @@ TRAIN:
 VALID:
     short_size: 256
     target_size: 224
-    num_reader_threads: 12
-    buf_size: 1024
-    batch_size: 128
+    num_reader_threads: 8
+    buf_size: 64
+    batch_size: 32
     filelist: "./data/dataset/ucf101/ucf101_val_split_1_videos.txt"
 
 TEST:
     short_size: 256
     target_size: 224
-    num_reader_threads: 12
-    buf_size: 1024
-    batch_size: 64
+    num_reader_threads: 8
+    buf_size: 64
+    batch_size: 32
     filelist: "./data/dataset/ucf101/ucf101_val_split_1_videos.txt"

dygraph/tsn/single_tsn_video.yaml

@@ -95,53 +95,6 @@ def parse_args():
     return args
 
 
-def decompress(path):
-    t = tarfile.open(path)
-    print("path[0] {}".format(os.path.split(path)[0]))
-    t.extractall(path=os.path.split(path)[0])
-    t.close()
-
-
-def download(url, path):
-    weight_dir = os.path.split(path)[0]
-    if not os.path.exists(weight_dir):
-        os.makedirs(weight_dir)
-
-    path = path + ".tar.gz"
-    print("path {}".format(path))
-    wget.download(url, path)
-    decompress(path)
-
-
-def pretrain_info():
-    return (
-        'ResNet50_pretrained',
-        'https://paddlemodels.bj.bcebos.com/video_classification/ResNet50_pretrained.tar.gz'
-    )
-
-
-def download_pretrained(pretrained):
-    if pretrained is not None:
-        WEIGHT_DIR = pretrained
-    else:
-        WEIGHT_DIR = os.path.join(os.path.expanduser('~'), '.paddle', 'weights')
-
-    path, url = pretrain_info()
-    if not path:
-        return None
-
-    path = os.path.join(WEIGHT_DIR, path)
-    if not os.path.isdir(WEIGHT_DIR):
-        logger.info('{} not exists, will be created automatically.'.format(
-            WEIGHT_DIR))
-        os.makedirs(WEIGHT_DIR)
-    if os.path.exists(path):
-        return path
-    logger.info("Download pretrain weights of ResNet50 from {}".format(url))
-    download(url, path)
-    return path
-
-
 def init_model(model, pre_state_dict):
     param_state_dict = {}
     model_dict = model.state_dict()
@@ -224,17 +177,14 @@ def train(args):
     train_config = merge_configs(config, 'train', vars(args))
     valid_config = merge_configs(config, 'valid', vars(args))
     print_configs(train_config, 'Train')
-
-    # get the pretrained weights
-    pretrained_path = download_pretrained(args.pretrain)
-
     use_data_parallel = args.use_data_parallel
+
     trainer_count = fluid.dygraph.parallel.Env().nranks
 
     # (data_parallel step1/6)
     place = fluid.CUDAPlace(fluid.dygraph.parallel.Env().dev_id) \
         if use_data_parallel else fluid.CUDAPlace(0)
-    pre_state_dict = fluid.load_program_state(pretrained_path)
+    pre_state_dict = fluid.load_program_state(args.pretrain)
 
     with fluid.dygraph.guard(place):
         if use_data_parallel:
@@ -342,7 +292,8 @@ def train(args):
                 model_path = os.path.join(
                     args.checkpoint,
                     "_" + model_path_pre + "_epoch{}".format(epoch))
-                fluid.dygraph.save_dygraph(video_model.state_dict(), model_path)
+                fluid.dygraph.save_dygraph(
+                    video_model.state_dict(), model_path)
                 fluid.dygraph.save_dygraph(optimizer.state_dict(), model_path)
 
             if args.validate: