Skip to content

Instantly share code, notes, and snippets.

@ak110
Last active December 28, 2017 13:47
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save ak110/8b9eed29f87fa5ead990c0264fd6c6ec to your computer and use it in GitHub Desktop.
Save ak110/8b9eed29f87fa5ead990c0264fd6c6ec to your computer and use it in GitHub Desktop.
TensorFlowで強制的にallow_growth = Trueをする ref: https://qiita.com/ak11/items/875c0f520ff1e231ee0c
import importlib.machinery
import sys
class CustomFinder(importlib.machinery.PathFinder):
def find_spec(self, fullname, path=None, target=None):
if fullname == 'tensorflow.python.client.session':
spec = super().find_spec(fullname, path, target)
loader = CustomLoader(fullname, spec.origin)
return importlib.machinery.ModuleSpec(fullname, loader)
return None
class CustomLoader(importlib.machinery.SourceFileLoader):
def exec_module(self, module):
r = super().exec_module(module)
self._patch(module)
return r
def _patch(self, module):
original_init = module.Session.__init__
def custom_init(self, *args, config=None, **kwargs):
import tensorflow as tf
if config is None:
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
return original_init(self, *args, config=config, **kwargs)
module.Session.__init__ = custom_init
return module
sys.meta_path.insert(0, CustomFinder())
$ python -c 'import tensorflow as tf; print(tf.Session()._config)'
(略)
gpu_options {
allow_growth: true
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment