提交 e7c7cbaa 编写于 作者: M minqiyang

Port new added files to Python3

上级 e4057d07
...@@ -14,12 +14,14 @@ ...@@ -14,12 +14,14 @@
""" """
This module privides a memory usage calculate function for user. This module privides a memory usage calculate function for user.
The purpose of this API is to allow users to estimate memory usage of 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 a program under a special batch size, then user can set appropriate
batch size to fully utilize a GPU. batch size to fully utilize a GPU.
This API is still under active development and may change drastically. This API is still under active development and may change drastically.
""" """
import six
from .. import core from .. import core
from ..framework import Program, Variable from ..framework import Program, Variable
...@@ -45,15 +47,15 @@ def memory_usage(program, batch_size): ...@@ -45,15 +47,15 @@ def memory_usage(program, batch_size):
Args: Args:
program(Program): The current Program. program(Program): The current Program.
batch_size(int): The current input data batch_size. batch_size(int): The current input data batch_size.
Returns: Returns:
min_total_memory(float): the estimate memory usage lower bound. min_total_memory(float): the estimate memory usage lower bound.
max_total_memory(float): the estimate memory usage upper bound. max_total_memory(float): the estimate memory usage upper bound.
unit_str(string): the unit of estimate usage result. unit_str(string): the unit of estimate usage result.
Examples: Examples:
>>> import paddle.fluid as fluid >>> import paddle.fluid as fluid
>>> lower_usage, upper_usage, unit = fluid.contrib.memory_usage( >>> lower_usage, upper_usage, unit = fluid.contrib.memory_usage(
fluid.default_main_program(), batch_size=10) fluid.default_main_program(), batch_size=10)
...@@ -72,7 +74,7 @@ def memory_usage(program, batch_size): ...@@ -72,7 +74,7 @@ def memory_usage(program, batch_size):
# Get the var_name list of first block and calculate # Get the var_name list of first block and calculate
total_memory = 0.0 total_memory = 0.0
for var in program.global_block().vars.itervalues(): for var in six.itervalues(program.global_block().vars):
data_count = 1 data_count = 1
for x in var.shape: for x in var.shape:
if x == -1: if x == -1:
...@@ -81,10 +83,10 @@ def memory_usage(program, batch_size): ...@@ -81,10 +83,10 @@ def memory_usage(program, batch_size):
data_count *= x data_count *= x
var_memory = data_count * dtype_to_size[var.dtype] var_memory = data_count * dtype_to_size[var.dtype]
if DEBUG: if DEBUG:
print "%s memory usage: %d" % (var.name, var_memory) print("%s memory usage: %d" % (var.name, var_memory))
total_memory += var_memory total_memory += var_memory
if DEBUG: if DEBUG:
print "total memory usage: %.2f" % (total_memory) print("total memory usage: %.2f" % (total_memory))
# Convert appropriate unit # Convert appropriate unit
unit_str = "B" unit_str = "B"
......
...@@ -160,6 +160,7 @@ def get_model(): ...@@ -160,6 +160,7 @@ def get_model():
avg_cost = transformer(use_feed=False) avg_cost = transformer(use_feed=False)
optimizer = fluid.optimizer.Adam() optimizer = fluid.optimizer.Adam()
optimizer.minimize(avg_cost) optimizer.minimize(avg_cost)
fluid.memory_optimize(fluid.default_main_program())
return avg_cost return avg_cost
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册