tensorflow 中 inter_op 和 intra

[root@localhost custom-resnet-v2]# cat runme.sh

#python demo_slim.py -h

#python demo_slim.py --cpu_num 8 --inter_op_threads 1 --intra_op_threads 8 --dump_timeline True

# export KMP_AFFINITY=verbose,granularity=fine,proclist=[0,1,2,3],explicit

# numactl -C 0-3 python demo_slim.py --cpu_num 4 --inter_op_threads 1 --intra_op_threads 4 >& run1.log &

export OMP_NUM_THREADS=8

python demo_slim.py --cpu_num 8 --inter_op_threads 1 --intra_op_threads 8

[root@localhost custom-resnet-v2]# cat demo_slim.py

# coding: utf8

import os

import sys

import numpy as np

import tensorflow as tf

from tensorflow.python.client import timeline

import argparse

import time

def make_fake_input(batch_size, input_height, input_width, input_channel):

        im = np.zeros((input_height,input_width,input_channel), np.uint8)

        im[:,:,:] = 1

        images = np.zeros((batch_size, input_height, input_width, input_channel), dtype=np.float32)

        for i in xrange(batch_size):

                images[i, 0:im.shape[0], 0:im.shape[1], :] = im

                #channel_swap = (0, 3, 1, 2)  # caffe

                #images = np.transpose(images, channel_swap)

                #cv2.imwrite("test.jpg", im)

        return images

def get_parser():

        """

        create a parser to parse argument "--cpu_num --inter_op_threads --intra_op_threads"

        """

        parser = argparse.ArgumentParser(description="Specify tensorflow parallelism")

        parser.add_argument("--cpu_num", dest="cpu_num", default=1, help="specify how many cpus to use.(default: 1)")

        parser.add_argument("--inter_op_threads", dest="inter_op_threads", default=1, help="specify max inter op parallelism.(default: 1)")

        parser.add_argument("--intra_op_threads", dest="intra_op_threads", default=1, help="specify max intra op parallelism.(default: 1)")

        parser.add_argument("--dump_timeline", dest="dump_timeline", default=False, help="specify to dump timeline.(default: False)")

        return parser

def main():

        parser = get_parser()

        args = parser.parse_args()

        #parser.print_help()

        cpu_num = int(args.cpu_num)

        inter_op_threads = int(args.inter_op_threads)

        intra_op_threads = int(args.intra_op_threads)

        dump_timeline = bool(args.dump_timeline)

        print("cpu_num: ", cpu_num)

        print("inter_op_threads: ", inter_op_threads)

        print("intra_op_threads: ", intra_op_threads)

        print("dump_timeline: ", dump_timeline)

        config = tf.ConfigProto(device_count={"CPU": cpu_num}, # limit to num_cpu_core CPU usage

                inter_op_parallelism_threads = inter_op_threads,

                intra_op_parallelism_threads = intra_op_threads,

                log_device_placement=False)

        with tf.Session(config = config) as sess:

                imgs = make_fake_input(1, 224, 224, 3)

                #init_start = time.time()

                saver = tf.train.import_meta_graph("slim_model/slim_model.ckpt.meta")

                saver.restore(sess, tf.train.latest_checkpoint("slim_model/"))

                ## Operations

                #for op in tf.get_default_graph().get_operations():

                #       print(op.name)

                #       print(op.values())

                graph = tf.get_default_graph()

                input_data = graph.get_tensor_by_name("Placeholder:0")

                fc6 = graph.get_tensor_by_name("resnet_v2/avg_fc_fc6_Conv2D/BiasAdd:0")

                #init_end = time.time()

                #print("initialization time: ", init_end-init_start, "s")

                time_start = time.time()

                for step in range(200):

                        if dump_timeline:

                                run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)

                                run_metadata = tf.RunMetadata()

                                result = sess.run(fc6, feed_dict={input_data:imgs}, options=run_options, run_metadata=run_metadata)

                                tm = timeline.Timeline(run_metadata.step_stats)

                                ctf = tm.generate_chrome_trace_format()

                                with open('timeline.json', 'w') as f:

                                        f.write(ctf)

                        else:

                                result = sess.run(fc6, feed_dict={input_data:imgs})

                        print(result[0][0][0])

                time_end = time.time()

                avg_time = (time_end-time_start) * 1000 / 200;

                print("AVG Time: ", avg_time, " ms")

        return 0

if __name__ == "__main__":

        sys.exit(main())

巴特西

tensorflow 中 inter_op 和 intra_op

最新文章

热门文章