(module)
| 42 | return '\n'.join(output) |
| 43 | |
| 44 | def model_size_estimation(module): |
| 45 | def estimate_size(model): |
| 46 | param_size = sum(p.nelement() * p.element_size() for p in model.parameters()) |
| 47 | buffer_size = sum(b.nelement() * b.element_size() for b in model.buffers()) |
| 48 | return param_size + buffer_size |
| 49 | |
| 50 | total_size = estimate_size(module) |
| 51 | output = [f'Total model size: {total_size / 1024**2:.2f} MB', '---------------------------'] |
| 52 | |
| 53 | for name, child in module.named_children(): |
| 54 | child_size = estimate_size(child) |
| 55 | output.append(f'{name} size: {child_size / 1024**2:.2f} MB') |
| 56 | |
| 57 | return '\n'.join(output) |
| 58 | |
| 59 | def layer_param_distribution(module): |
| 60 | def count_parameters(model): |
nothing calls this directly
no test coverage detected