| 351 | return GradRejust.apply(x, grad_scale) |
| 352 | |
| 353 | def count_parameters_num(model): |
| 354 | count = 0 |
| 355 | count_fc = 0 |
| 356 | param_dict = {name:param for name,param in model.named_parameters()} |
| 357 | param_keys = param_dict.keys() |
| 358 | for m_name, m in model.named_modules(): |
| 359 | if isinstance(m, nn.Conv2d) or isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.SyncBatchNorm): |
| 360 | weight_name = m_name + '.weight' |
| 361 | bias_name = m_name + '.bias' |
| 362 | if weight_name in param_keys: |
| 363 | temp_params = param_dict[weight_name] |
| 364 | count += temp_params.data.nelement() |
| 365 | if bias_name in param_keys: |
| 366 | temp_params = param_dict[bias_name] |
| 367 | count += temp_params.data.nelement() |
| 368 | elif isinstance(m, nn.Linear): |
| 369 | weight_name = m_name + '.weight' |
| 370 | bias_name = m_name + '.bias' |
| 371 | if weight_name in param_keys: |
| 372 | temp_params = param_dict[weight_name] |
| 373 | count_fc += temp_params.data.nelement() |
| 374 | if bias_name in param_keys: |
| 375 | temp_params = param_dict[bias_name] |
| 376 | count_fc += temp_params.data.nelement() |
| 377 | sync_print('Number of conv/bn params: %.2fM' % (count / 1e6)) |
| 378 | sync_print('Number of linear params: %.2fM' % (count_fc / 1e6)) |
| 379 | |
| 380 | def get_gpu_memory_map(): |
| 381 | """Get the current gpu usage. |