Compute the number of parameters. Args: model (model): model to count the number of parameters.
(model, ignore_bn=False)
| 34 | |
| 35 | |
| 36 | def params_count(model, ignore_bn=False): |
| 37 | """ |
| 38 | Compute the number of parameters. |
| 39 | Args: |
| 40 | model (model): model to count the number of parameters. |
| 41 | """ |
| 42 | if not ignore_bn: |
| 43 | return np.sum([p.numel() for p in model.parameters()]).item() |
| 44 | else: |
| 45 | count = 0 |
| 46 | for m in model.modules(): |
| 47 | if not isinstance(m, nn.BatchNorm3d): |
| 48 | for p in m.parameters(recurse=False): |
| 49 | count += p.numel() |
| 50 | return count |
| 51 | |
| 52 | |
| 53 | def gpu_mem_usage(): |