diff --git a/paddleslim/analysis/__init__.py b/paddleslim/analysis/__init__.py index 9d0531501ca43921438ee5b2fb58ac0ad2396d1b..9caa0d24006a3e59f2d39c646d247b7e68480f96 100644 --- a/paddleslim/analysis/__init__.py +++ b/paddleslim/analysis/__init__.py @@ -11,3 +11,10 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import flops as flops_module +from flops import * +import model_size as model_size_module +from model_size import * +__all__ = [] +__all__ += flops_module.__all__ +__all__ += model_size_module.__all__ diff --git a/paddleslim/analysis/model_size.py b/paddleslim/analysis/model_size.py new file mode 100644 index 0000000000000000000000000000000000000000..34574d5d53764810185112d7122aeb5b99d74682 --- /dev/null +++ b/paddleslim/analysis/model_size.py @@ -0,0 +1,31 @@ +# Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License" +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +from ..core import GraphWrapper + +__all__ = ["model_size"] + + +def model_size(program): + """ + Get total value numbers of all parameters. + Args: + program(Program): The program used to calculate model size. + """ + size = 0 + for block in program.blocks: + for param in block.all_parameters(): + size += np.product(param.shape) + return size diff --git a/tests/test_model_size.py b/tests/test_model_size.py new file mode 100644 index 0000000000000000000000000000000000000000..314450eb449507a971aa3827a8d841e88e10a69e --- /dev/null +++ b/tests/test_model_size.py @@ -0,0 +1,40 @@ +# Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License" +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import sys +sys.path.append("../") +import unittest +import paddle.fluid as fluid +from paddleslim.analysis import model_size +from layers import conv_bn_layer + + +class TestPrune(unittest.TestCase): + def test_prune(self): + main_program = fluid.Program() + startup_program = fluid.Program() + with fluid.program_guard(main_program, startup_program): + input = fluid.data(name="image", shape=[None, 3, 16, 16]) + conv1 = conv_bn_layer(input, 8, 3, "conv1") + conv2 = conv_bn_layer(conv1, 8, 3, "conv2") + sum1 = conv1 + conv2 + conv3 = conv_bn_layer(sum1, 8, 3, "conv3") + conv4 = conv_bn_layer(conv3, 8, 3, "conv4") + sum2 = conv4 + sum1 + conv5 = conv_bn_layer(sum2, 8, 3, "conv5") + conv6 = conv_bn_layer(conv5, 8, 3, "conv6") + self.assertTrue(3288 == model_size(main_program)) + + +if __name__ == '__main__': + unittest.main()