tf_sess_config.py 1.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445
  1. # Copyright (c) 2021
  2. #
  3. # This file is part of FFmpeg.
  4. #
  5. # FFmpeg is free software; you can redistribute it and/or
  6. # modify it under the terms of the GNU Lesser General Public
  7. # License as published by the Free Software Foundation; either
  8. # version 2.1 of the License, or (at your option) any later version.
  9. #
  10. # FFmpeg is distributed in the hope that it will be useful,
  11. # but WITHOUT ANY WARRANTY; without even the implied warranty of
  12. # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
  13. # Lesser General Public License for more details.
  14. #
  15. # You should have received a copy of the GNU Lesser General Public
  16. # License along with FFmpeg; if not, write to the Free Software
  17. # Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  18. # ==============================================================================
  19. # verified with Python 3.6.8 on CentOS 7.2
  20. import tensorflow as tf
  21. visible_device_list = '0' # use , separator for more GPUs like '0, 1'
  22. per_process_gpu_memory_fraction = 0.9 # avoid out of memory
  23. intra_op_parallelism_threads = 2 # default in tensorflow
  24. inter_op_parallelism_threads = 5 # default in tensorflow
  25. gpu_options = tf.compat.v1.GPUOptions(
  26. per_process_gpu_memory_fraction = per_process_gpu_memory_fraction,
  27. visible_device_list = visible_device_list,
  28. allow_growth = True)
  29. config = tf.compat.v1.ConfigProto(
  30. allow_soft_placement = True,
  31. log_device_placement = False,
  32. intra_op_parallelism_threads = intra_op_parallelism_threads,
  33. inter_op_parallelism_threads = inter_op_parallelism_threads,
  34. gpu_options = gpu_options)
  35. s = config.SerializeToString()
  36. # print(list(map(hex, s))) # print by json if need
  37. print('a serialized protobuf string for TF_SetConfig, note the byte order is in normal order.')
  38. b = ''.join(format(b,'02x') for b in s)
  39. print('0x%s' % b) # print by hex format