diff --git a/ppcls/modeling/architectures/vgg.py b/ppcls/modeling/architectures/vgg.py index ccb94a9fccca7a967b653f26d764ed747322212b..06ba8d08c323da0a1c1d6bcfcca104817085885d 100644 --- a/ppcls/modeling/architectures/vgg.py +++ b/ppcls/modeling/architectures/vgg.py @@ -113,7 +113,8 @@ class VGGNet(nn.Layer): x = self._conv_block_4(x) x = self._conv_block_5(x) - x = paddle.reshape(x, [-1, x.shape[1]*x.shape[2]*x.shape[3]]) + _, c, h, w = list(x.shape) + x = paddle.reshape(x, [-1, c * h * w]) x = self._fc1(x) x = F.relu(x) x = self._drop(x) diff --git a/tools/infer/utils.py b/tools/infer/utils.py index f299189047eeca402745ed7153132bc362326570..9ea5e35422c57db49391a2e0551376068929e782 100644 --- a/tools/infer/utils.py +++ b/tools/infer/utils.py @@ -41,6 +41,7 @@ def parse_args(): parser.add_argument("--ir_optim", type=str2bool, default=True) parser.add_argument("--use_tensorrt", type=str2bool, default=False) parser.add_argument("--gpu_mem", type=int, default=8000) + parser.add_argument("--enable_profile", type=str2bool, default=False) parser.add_argument("--enable_benchmark", type=str2bool, default=False) parser.add_argument("--top_k", type=int, default=1) parser.add_argument("--enable_mkldnn", type=str2bool, default=False) @@ -81,6 +82,8 @@ def create_paddle_predictor(args): config.enable_mkldnn() config.set_cpu_math_library_num_threads(args.cpu_num_threads) + if args.enable_profile: + config.enable_profile() config.disable_glog_info() config.switch_ir_optim(args.ir_optim) # default true if args.use_tensorrt: