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