(model)
| 158 | return transform |
| 159 | |
| 160 | def get_parameter_number(model): |
| 161 | trainable_params, all_param = 0, 0 |
| 162 | for param in model.parameters(): |
| 163 | num_params = param.numel() |
| 164 | # if using DS Zero 3 and the weights are initialized empty |
| 165 | if num_params == 0 and hasattr(param, "ds_numel"): |
| 166 | num_params = param.ds_numel |
| 167 | |
| 168 | all_param += num_params |
| 169 | if param.requires_grad: |
| 170 | trainable_params += num_params |
| 171 | |
| 172 | return {'Total': all_param, 'Trainable': trainable_params} |
| 173 | |
| 174 | |
| 175 | local_rank = 0 |