TensorFlow uses recomputer technology

manpreet Best Answer 2 years ago

In order to understand how much GPU memory resources TensorFlow can economize, I did some experiments. My network is vgg-16 and dataset is mnist, batch size is 128. I estimate my GPU memory usage is 2.8GB. I close all the RewriterConfigs like:

rewrite_options.memory_optimization = rewriter_config_pb2.RewriterConfig.NO_MEM_OPT
rewrite_options.layout_optimizer = rewriter_config_pb2.RewriterConfig.OFF
rewrite_options.function_optimization = rewriter_config_pb2.RewriterConfig.OFF
rewrite_options.constant_folding = rewriter_config_pb2.RewriterConfig.OFF
rewrite_options.arithmetic_optimization = rewriter_config_pb2.RewriterConfig.OFF
rewrite_options.loop_optimization = rewriter_config_pb2.RewriterConfig.OFF
rewrite_options.dependency_optimization = rewriter_config_pb2.RewriterConfig.OFF

And I gradually reduce the amount of GPU memory usage.

I observed that TensorFlow's method is recomputer in lack of GPU memory. Why use recomputer technology instead of swapping technology? Is recompute faster than swapping?

In the memory_optimizer.cc:1226, it has the RecomputationRewritingPass function. After I close the function to test, I find that TensorFlow has another block do the same event. Why does TensorFlow have two blocks to do the same event?

Thanks.