Count total, classifier, and non-classifier trainable parameters in a PyTorch model. Returns values in terms of K (1000) parameters. Args: model: PyTorch model verbose: If True, print parameter counts for each layer Returns: Dictionary containin
(model: nn.Module, verbose: bool = False)
| 9 | return str(num) |
| 10 | |
| 11 | def count_parameters(model: nn.Module, verbose: bool = False) -> Dict[str, float]: |
| 12 | """ |
| 13 | Count total, classifier, and non-classifier trainable parameters in a PyTorch model. |
| 14 | Returns values in terms of K (1000) parameters. |
| 15 | |
| 16 | Args: |
| 17 | model: PyTorch model |
| 18 | verbose: If True, print parameter counts for each layer |
| 19 | |
| 20 | Returns: |
| 21 | Dictionary containing parameter counts in K (1000s) |
| 22 | """ |
| 23 | def is_classifier_layer(name: str) -> bool: |
| 24 | """Check if the layer is part of classifier based on common naming patterns""" |
| 25 | classifier_keywords = ['classifier', 'fc', 'linear', 'head'] |
| 26 | return any(keyword in name.lower() for keyword in classifier_keywords) |
| 27 | |
| 28 | total_params = 0 |
| 29 | classifier_params = 0 |
| 30 | non_classifier_params = 0 |
| 31 | |
| 32 | # Iterate through all parameters |
| 33 | for name, parameter in model.named_parameters(): |
| 34 | if parameter.requires_grad: |
| 35 | param_count = parameter.numel() |
| 36 | total_params += param_count |
| 37 | |
| 38 | if is_classifier_layer(name): |
| 39 | classifier_params += param_count |
| 40 | else: |
| 41 | non_classifier_params += param_count |
| 42 | |
| 43 | if verbose: |
| 44 | print(f"{name}: {format_params(param_count)} parameters " |
| 45 | f"{'(Classifier)' if is_classifier_layer(name) else '(Non-classifier)'}") |
| 46 | |
| 47 | results = { |
| 48 | 'total_trainable_params': total_params / 1000, # Convert to K |
| 49 | 'classifier_params': classifier_params / 1000, # Convert to K |
| 50 | 'non_classifier_params': non_classifier_params / 1000 # Convert to K |
| 51 | } |
| 52 | |
| 53 | print("\nSummary:") |
| 54 | print(f"Total trainable parameters (K): {format_params(total_params)}") |
| 55 | print(f"Classifier parameters (K): {format_params(classifier_params)}") |
| 56 | print(f"Non-classifier parameters (K): {format_params(non_classifier_params)}") |
| 57 | print(f"Classifier parameters percentage (K): {(classifier_params/total_params)*100:.2f}%") |
| 58 | |
| 59 | return results |
no test coverage detected