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.")
parser.add_argument('--config', help='config file')
parser.add_argument('--meta', help='metagraph file')
parser.add_argument(dest='model')
parser.add_argument(dest='output')
args = parser.parse_args()
assert args.config or args.meta, "Either config or metagraph must be present!"
with tf.Graph().as_default() as G:
if args.config:
logger.warn(
"Using a config script is not reliable. Please use metagraph.")
MODEL = imp.load_source('config_script', args.config).Model
M = MODEL()
with TowerContext('', is_training=False):
input = PlaceholderInput()
input.setup(M.get_inputs_desc())
M.build_graph(input)
else:
tf.train.import_meta_graph(args.meta)
# loading...
init = get_model_loader(args.model)
sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
init.init(sess)
# dump ...
with sess.as_default():
if args.output.endswith('npy') or args.output.endswith('npz'):
if args.gpu:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
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)
if args.gpu:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
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:
if args.gpu:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
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)
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)
parser = argparse.ArgumentParser()
parser.add_argument('--config', help='config file')
parser.add_argument('--meta', help='metagraph file')
parser.add_argument(dest='model')
parser.add_argument(dest='output')
args = parser.parse_args()
assert args.config or args.meta, "Either config or metagraph must be present!"
with tf.Graph().as_default() as G:
if args.config:
logger.warn("Using a config script is not reliable. Please use metagraph.")
MODEL = imp.load_source('config_script', args.config).Model
M = MODEL()
with TowerContext('', is_training=False):
input = PlaceholderInput()
input.setup(M.get_inputs_desc())
M.build_graph(*input.get_input_tensors())
else:
tf.train.import_meta_graph(args.meta)
# loading...
init = get_model_loader(args.model)
sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
init.init(sess)
# dump ...
with sess.as_default():
if args.output.endswith('npy') or args.output.endswith('npz'):
finalize_configs(is_training=False)
MODEL = PredModel(1, cfg.PREPROC.INPUT_SHAPE_EVAL)
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)