Compute statistics for the current model given the config. Args: model (model): model to perform analysis. cfg (CfgNode): configs. Details can be found in slowfast/config/defaults.py mode (str): Options include `flop` or `activation`. Compute either flop
(model, cfg, mode, use_train_input)
| 135 | |
| 136 | |
| 137 | def get_model_stats(model, cfg, mode, use_train_input): |
| 138 | """ |
| 139 | Compute statistics for the current model given the config. |
| 140 | Args: |
| 141 | model (model): model to perform analysis. |
| 142 | cfg (CfgNode): configs. Details can be found in |
| 143 | slowfast/config/defaults.py |
| 144 | mode (str): Options include `flop` or `activation`. Compute either flop |
| 145 | (gflops) or activation count (mega). |
| 146 | use_train_input (bool): if True, compute statistics for training. Otherwise, |
| 147 | compute statistics for testing. |
| 148 | |
| 149 | Returns: |
| 150 | float: the total number of count of the given model. |
| 151 | """ |
| 152 | assert mode in [ |
| 153 | "flop", |
| 154 | "activation", |
| 155 | ], "'{}' not supported for model analysis".format(mode) |
| 156 | if mode == "flop": |
| 157 | model_stats_fun = flop_count |
| 158 | elif mode == "activation": |
| 159 | model_stats_fun = activation_count |
| 160 | |
| 161 | # Set model to evaluation mode for analysis. |
| 162 | # Evaluation mode can avoid getting stuck with sync batchnorm. |
| 163 | model_mode = model.training |
| 164 | model.eval() |
| 165 | inputs = _get_model_analysis_input(cfg, use_train_input) |
| 166 | count_dict, *_ = model_stats_fun(model, inputs) |
| 167 | count = sum(count_dict.values()) |
| 168 | model.train(model_mode) |
| 169 | return count |
| 170 | |
| 171 | |
| 172 | def log_model_info(model, cfg, use_train_input=True): |
no test coverage detected