Port new added files to Python3

revert-12469-sum_op_dim_fix
minqiyang 7 years ago
parent e4057d071b
commit e7c7cbaa42

@ -20,6 +20,8 @@ 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
@ -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):
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():
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

Loading…
Cancel
Save