tools/python/tf_sess_config.py

   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
  20 # verified with Python 3.6.8 on CentOS 7.2
  21 import tensorflow as tf
  22
  23 visible_device_list = '0' # use , separator for more GPUs like '0, 1'
  24 per_process_gpu_memory_fraction = 0.9 # avoid out of memory
  25 intra_op_parallelism_threads = 2  # default in tensorflow
  26 inter_op_parallelism_threads = 5  # default in tensorflow
  27
  28 gpu_options = tf.compat.v1.GPUOptions(
  29               per_process_gpu_memory_fraction = per_process_gpu_memory_fraction,
  30               visible_device_list = visible_device_list,
  31               allow_growth = True)
  32
  33 config = tf.compat.v1.ConfigProto(
  34          allow_soft_placement = True,
  35          log_device_placement = False,
  36          intra_op_parallelism_threads = intra_op_parallelism_threads,
  37          inter_op_parallelism_threads = inter_op_parallelism_threads,
  38          gpu_options = gpu_options)
  39
  40 s = config.SerializeToString()
  41 # print(list(map(hex, s)))  # print by json if need
  42
  43 print('a serialized protobuf string for TF_SetConfig, note the byte order is in normal order.')
  44 b = ''.join(format(b,'02x') for b in s)
  45 print('0x%s' % b) # print by hex format