def test_print_shape(self): model = create_model2() input = torch.zeros(model.OP_INPUT) model_str_before = str(model) fest = flops_utils.FlopsEstimation(model) model_str_after = str(model) # make sure str was not affected when it is not enabled self.assertEqual(model_str_before, model_str_after) with fest.enable(): model(input) fest.add_flops_info() model_str = str(model) print(model_str) GT_SHAPES_STRS = [ "input_shapes=[[2, 3, 16, 16]], output_shapes=[2, 8, 16, 16]", "input_shapes=[[2, 8, 8, 8]], output_shapes=[2, 4, 8, 8]", "ReLU(input_shapes=[[2, 4, 8, 8]], output_shapes=[2, 4, 8, 8], nparams=0.0, nflops=0.0)", # noqa "nparams=0.00108, nflops=0.221184", ] for x in GT_SHAPES_STRS: self.assertIn(x, model_str) # make sure the additional informaiton are cleaned up model_str_clean = str(model) GT_SHAPES = ["input_shapes", "output_shapes", "nparams", "nflops"] for x in GT_SHAPES: self.assertNotIn(x, model_str_clean)
def test_duplicated(self): """Make sure handles subclasses propertly for mock""" model = nn.Sequential(M1(), M2(), nn.Conv2d(3, 4, 3), nn.ConvTranspose2d(4, 4, 3), M1()) input = torch.zeros([1, 3, 4, 4]) fest = flops_utils.FlopsEstimation(model) count = 0 flops = [] def flops_callback(fest, model, model_data): nonlocal count nparams, nflops = fest.get_flops() flops.append({"nparams": nparams, "nflops": nflops}) if count >= 2: fest.set_enable(False) count += 1 fest.set_callback(flops_callback) fest.set_enable(True) for _ in range(5): model(input) gt_flops = [{"nparams": 0.000252, "nflops": 0.002736}] * 3 self.assertEqual(gt_flops, flops)
def test_callback(self): model = create_model2() input = torch.zeros(model.OP_INPUT) fest = flops_utils.FlopsEstimation(model) count = 0 flops = [] def flops_callback(fest, model, model_data): nonlocal count nparams, nflops = fest.get_flops() flops.append({"nparams": nparams, "nflops": nflops}) if count >= 2: fest.set_enable(False) count += 1 fest.set_callback(flops_callback) fest.set_enable(True) for _ in range(5): model(input) gt_flops = [{"nparams": 0.00108, "nflops": 0.221184}] * 3 self.assertEqual(gt_flops, flops)
def print_flops(model, first_batch): logger.info("Evaluating model's number of parameters and FLOPS") model_flops = copy.deepcopy(model) model_flops.eval() fest = flops_utils.FlopsEstimation(model_flops) with fest.enable(): model_flops(first_batch) fest.add_flops_info() model_str = str(model_flops) logger.info(model_str) return model_str
def test_get_flops(self): model = create_model2() fest = flops_utils.FlopsEstimation(model) input = torch.zeros(model.OP_INPUT) with fest.enable(): model(input) nparams, nflops = fest.get_flops() self.assertAlmostEqual(nparams, 0.00108) self.assertAlmostEqual(nflops, 0.221184)
def dump_flops_info(model, inputs, output_dir): """ Dump flops information about model, using the given model inputs. Information are dumped to output_dir using various flop counting tools in different formats. Only a simple table is printed to terminal. """ if not comm.is_main_process(): return logger.info("Evaluating model's number of parameters and FLOPS") model = copy.deepcopy(model) model.eval() # 1. using mobile_cv flop counter try: fest = flops_utils.FlopsEstimation(model) with fest.enable(): model(inputs) fest.add_flops_info() model_str = str(model) output_file = os.path.join(output_dir, "flops_str_mobilecv.txt") with PathManager.open(output_file, "w") as f: f.write(model_str) logger.info(f"Flops info written to {output_file}") except Exception: logger.exception( "Failed to estimate flops using mobile_cv's FlopsEstimation") # 2. using d2/fvcore's flop counter try: flops = FlopCountAnalysis(model, inputs) # 2.1: dump as model str model_str = flop_count_str(flops) output_file = os.path.join(output_dir, "flops_str_fvcore.txt") with PathManager.open(output_file, "w") as f: f.write(model_str) logger.info(f"Flops info written to {output_file}") # 2.2: dump as table flops_table = flop_count_table(flops, max_depth=10) output_file = os.path.join(output_dir, "flops_table_fvcore.txt") with PathManager.open(output_file, "w") as f: f.write(flops_table) logger.info(f"Flops table written to {output_file}") # 2.3: print a table with a shallow depth flops_table = flop_count_table(flops, max_depth=3) logger.info("Flops table:\n" + flops_table) except Exception: logger.exception( "Failed to estimate flops using detectron2's FlopCountAnalysis")
def add_print_flops_callback(cfg, model, disable_after_callback=True): def _print_flops_callback(self, model, model_data): self.add_flops_info() logger.info("Callback: model flops info:\n{}".format(model)) def _guess_batch_size(): # Inputs are meta-arch dependent, the most general solution will be # adding a function like `get_batch_size()` to each meta arch ret = 1 try: model_input_shapes = model_data(model)["input_shapes"] assert isinstance(model_input_shapes, list) assert len(model_input_shapes) > 0 # assuming the first input is a list of images ret = len(model_input_shapes[0]) except Exception: ret = cfg.SOLVER.IMS_PER_BATCH // comm.get_world_size() logger.warning( "Could not get batch size, compute from" f" `cfg.SOLVER.IMS_PER_BATCH`={ret}" ) pass return ret nparams, nflops = self.get_flops() batch_size = _guess_batch_size() nflops_single = nflops / batch_size logger.info( f"Model parameters (M): {nparams}, " f"MFlops (batch_size={batch_size}): {nflops}, " f"MFlops (batch_size=1): {nflops_single}" ) if disable_after_callback: self.set_enable(False) fest = flops_utils.FlopsEstimation(model).set_callback(_print_flops_callback) logger.info("Added callback to log flops info after the first inference") fest.set_enable(True) return fest
def dump_flops_info(model, inputs, output_dir, use_eval_mode=True): """ Dump flops information about model, using the given model inputs. Information are dumped to output_dir using various flop counting tools in different formats. Only a simple table is printed to terminal. Args: inputs: a tuple of positional arguments used to call model with. use_eval_mode: turn the model into eval mode for flop counting. Otherwise, will use the original mode. It's recommended to use eval mode, because training mode typically follows a different codepath. """ if not comm.is_main_process(): return logger.info("Evaluating model's number of parameters and FLOPS") try: model = copy.deepcopy(model) except Exception: logger.info("Failed to deepcopy the model and skip FlopsEstimation.") return # delete other forward_pre_hooks so they are not simultaneously called for k in model._forward_pre_hooks: del model._forward_pre_hooks[k] if use_eval_mode: model.eval() inputs = copy.deepcopy(inputs) # 1. using mobile_cv flop counter try: fest = flops_utils.FlopsEstimation(model) with fest.enable(): model(*inputs) fest.add_flops_info() model_str = str(model) output_file = os.path.join(output_dir, "flops_str_mobilecv.txt") with PathManager.open(output_file, "w") as f: f.write(model_str) logger.info(f"Flops info written to {output_file}") except Exception: logger.exception( "Failed to estimate flops using mobile_cv's FlopsEstimation") # 2. using d2/fvcore's flop counter output_file = os.path.join(output_dir, "flops_str_fvcore.txt") try: flops = FlopCountAnalysis(model, inputs) # 2.1: dump as model str model_str = flop_count_str(flops) with PathManager.open(output_file, "w") as f: f.write(model_str) logger.info(f"Flops info written to {output_file}") # 2.2: dump as table flops_table = flop_count_table(flops, max_depth=10) output_file = os.path.join(output_dir, "flops_table_fvcore.txt") with PathManager.open(output_file, "w") as f: f.write(flops_table) logger.info(f"Flops table (full version) written to {output_file}") # 2.3: print a table with a shallow depth flops_table = flop_count_table(flops, max_depth=3) logger.info("Flops table:\n" + flops_table) except Exception: with PathManager.open(output_file, "w") as f: traceback.print_exc(file=f) logger.warning( "Failed to estimate flops using detectron2's FlopCountAnalysis. " f"Error written to {output_file}.") flops = float("nan") return flops