def calc_flops(model): # manually build the graph with batch=1 input_desc = [ InputDesc(tf.float32, [1, 224, 224, 3], 'input'), InputDesc(tf.int32, [1], 'label') ] input = PlaceholderInput() input.setup(input_desc) with TowerContext('', is_training=False): model.build_graph(*input.get_input_tensors()) model_utils.describe_trainable_vars() tf.profiler.profile( tf.get_default_graph(), cmd='op', options=tf.profiler.ProfileOptionBuilder.float_operation()) logger.info( "Note that TensorFlow counts flops in a different way from the paper.") logger.info( "TensorFlow counts multiply+add as two flops, however the paper counts them " "as 1 flop because it can be executed in one instruction.")
model = Model() if args.eval: batch = 128 # something that can run on one gpu ds = get_data('val', batch) eval_on_ILSVRC12(model, get_model_loader(args.load), ds) elif args.flops: # manually build the graph with batch=1 input_desc = [ InputDesc(tf.float32, [1, 224, 224, 3], 'input'), InputDesc(tf.int32, [1], 'label') ] input = PlaceholderInput() input.setup(input_desc) with TowerContext('', is_training=True): model.build_graph(*input.get_input_tensors()) tf.profiler.profile( tf.get_default_graph(), cmd='op', options=tf.profiler.ProfileOptionBuilder.float_operation()) else: logger.set_logger_dir(os.path.join('train_log', 'shufflenet')) nr_tower = max(get_nr_gpu(), 1) config = get_config(model, nr_tower) if args.load: config.session_init = get_model_loader(args.load) launch_train_with_config(config, SyncMultiGPUTrainerParameterServer(nr_tower))
model = Model() if args.eval: batch = 128 # something that can run on one gpu ds = get_data('val', batch) eval_on_ILSVRC12(model, get_model_loader(args.load), ds) elif args.flops: # manually build the graph with batch=1 input_desc = [ InputDesc(tf.float32, [1, 224, 224, 3], 'input'), InputDesc(tf.int32, [1], 'label') ] input = PlaceholderInput() input.setup(input_desc) with TowerContext('', is_training=True): model.build_graph(*input.get_input_tensors()) tf.profiler.profile( tf.get_default_graph(), cmd='op', options=tf.profiler.ProfileOptionBuilder.float_operation()) else: logger.set_logger_dir( os.path.join('train_log', 'shufflenet')) nr_tower = max(get_nr_gpu(), 1) config = get_config(model, nr_tower) if args.load: config.session_init = get_model_loader(args.load) launch_train_with_config(config, SyncMultiGPUTrainerParameterServer(nr_tower))
in_shape = [ 1, cfg.PREPROC.INPUT_SHAPE_EVAL[0], cfg.PREPROC.INPUT_SHAPE_EVAL[1], 3 ] pred_config = PredictConfig( model=MODEL, input_names=MODEL.get_inference_tensor_names()[0], output_names=MODEL.get_inference_tensor_names()[1]) pred_config.inputs_desc[0] = InputDesc(type=tf.float32, shape=in_shape, name='data') inputs = PlaceholderInput() inputs.setup(pred_config.inputs_desc) with PredictTowerContext(''): MODEL.build_graph(*inputs.get_input_tensors()) model_utils.describe_trainable_vars() tf.profiler.profile( tf.get_default_graph(), cmd='op', options=tf.profiler.ProfileOptionBuilder.float_operation()) sys.exit(0) assert args.load # is_training = False finalize_configs(is_training=False) # define model input_shape = cfg.PREPROC.INPUT_SHAPE_EVAL batch_size = cfg.PREPROC.BATCH_SIZE