demo_yolo_v2.py 文件源码

def __init__(self):
        model_name = 'yolov2-coco'
        model_dir = './models'
        gpu_id = 4
        self.gpu_utility = 0.9

        self.pb_file = '{}/{}.pb'.format(model_dir, model_name)
        self.meta_file = '{}/{}.meta'.format(model_dir, model_name)
        self.batch = 4

        self.graph = tf.Graph()
        with tf.device('/gpu:1'):
            with self.graph.as_default() as g:
                self.build_from_pb()
                gpu_options = tf.GPUOptions(allow_growth=True)
                sess_config = tf.ConfigProto(gpu_options=gpu_options, log_device_placement=False)
                self.sess = tf.Session(config = sess_config)
                self.sess.run(tf.global_variables_initializer())
        return