diff --git a/python/paddle/fluid/contrib/__init__.py b/python/paddle/fluid/contrib/__init__.py index 12cd5d918e93181c6b7e328e6aee4ad941b0a0da..9de9e9504510baec9aefb47f91793c364450795a 100644 --- a/python/paddle/fluid/contrib/__init__.py +++ b/python/paddle/fluid/contrib/__init__.py @@ -14,5 +14,7 @@ import decoder from decoder import * +import memory_usage_calc +from memory_usage_calc import * -__all__ = decoder.__all__ +__all__ = decoder.__all__ + memory_usage_calc.__all__ diff --git a/python/paddle/fluid/contrib/memory_usage_calc.py b/python/paddle/fluid/contrib/memory_usage_calc.py new file mode 100644 index 0000000000000000000000000000000000000000..7b9c7afee03e529581934de46b8c3c36386c7777 --- /dev/null +++ b/python/paddle/fluid/contrib/memory_usage_calc.py @@ -0,0 +1,99 @@ +# Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserve. +# +# 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. +""" +This module privides a memory usage calculate function for user. +The purpose of this API is to allow users to estimate memory usage of +a program under a special batch size, then user can set appropriate +batch size to fully utilize a GPU. + +This API is still under active development and may change drastically. +""" + +from .. import core +from ..framework import Program, Variable + +__all__ = ['MemoryInfo'] + +DEBUG = False + +dtype_to_size = { + core.VarDesc.VarType.FP16: 2, + core.VarDesc.VarType.FP32: 4, + core.VarDesc.VarType.FP64: 8, + core.VarDesc.VarType.INT16: 2, + core.VarDesc.VarType.INT32: 4, + core.VarDesc.VarType.INT64: 8, + core.VarDesc.VarType.BOOL: 1, + core.VarDesc.VarType.UINT8: 1, +} + + +class MemoryInfo(object): + def __init__(self, program): + if not isinstance(program, Program): + raise TypeError( + "Calculating Memory Usage requires Program as its Parameter." + "But you passed in %s" % (type(prgram))) + self._program = program + + def _has_var(self, block, var_name): + return block.has_var(str(var_name)) + + def _find_var(self, block, var_name): + return block.var(str(var_name)) + + def get_memory_usage(self, batch_size, with_details=False): + + # get the first block of program + first_block = self._program.global_block() + + # get the var_name list of first block + # TODO(chenweihang): not find the API get block's var list directly + total_memory = 0.0 + for var in self._program.list_vars(): + if DEBUG: + print "All Block's Var: %s" % (var.name) + # TODO(chenweihang): why not used program.list_vars() + # calculate all variable's memory directly? + if self._has_var(first_block, var.name): + if DEBUG: + print "First Block's Var: %s" % (var.name) + print "Var's shape: ", var.shape + print "Var's dtype: ", var.dtype + data_count = 1 + for x in var.shape: + if x == -1: + data_count *= batch_size + else: + data_count *= x + var_memory = data_count * dtype_to_size[var.dtype] + if DEBUG: + print "Var's memory: %d" % (var_memory) + total_memory += var_memory + + # Convert unit and make result string + result_str = "- With current batch size, memory usage is about " + unit_str = " B." + if total_memory > 1024: + total_memory /= 1024 + unit_str = " KB." + if total_memory > 1024: + total_memory /= 1024 + unit_str = " MB." + + # Append extra memory consumption (5% - 10%) + result_str += str(round(total_memory * 1.05, 3)) + " - " \ + + str(round(total_memory * 1.10, 3)) + unit_str + + return result_str