def log_model_info(model, writer_eval=None):
"""Logs model info"""
logger.info('Model:\n{}'.format(model))
params = mu.params_count(model)
flops = mu.flops_count(model)
logger.info('Params: {:,}'.format(params))
logger.info('Flops: {:,}'.format(flops))
logger.info('Number of node: {:,}'.format(cfg.RGRAPH.GROUP_NUM))
# logger.info('{}, {}'.format(params,flops))
if writer_eval is not None:
writer_eval.add_scalar('Params', params, 1)
writer_eval.add_scalar('Flops', flops, 1)
return params, flops
def build_model_stats(mode='flops'):
"""Builds the model."""
assert cfg.MODEL.TYPE in _MODEL_TYPES.keys(), \
'Model type \'{}\' not supported'.format(cfg.MODEL.TYPE)
assert cfg.NUM_GPUS <= torch.cuda.device_count(), \
'Cannot use more GPU devices than available'
# Construct the model
model = _MODEL_TYPES[cfg.MODEL.TYPE]()
if mode == 'flops':
flops = mu.flops_count(model)
return flops
else:
params = mu.params_count(model)
return params
def log_model_info(model):
"""Logs model info"""
logger.info("Model:\n{}".format(model))
logger.info("Params: {:,}".format(mu.params_count(model)))
logger.info("Flops: {:,}".format(mu.flops_count(model)))
def log_model_info(model):
"""Logs model info"""
logger.info('Model:\n{}'.format(model))
logger.info('Params: {:,}'.format(mu.params_count(model)))
logger.info('Flops: {:,}'.format(mu.flops_count(model)))