(m)
| 248 | custom_ops = {} |
| 249 | |
| 250 | def add_hooks(m): |
| 251 | if len(list(m.children())) > 0: |
| 252 | return |
| 253 | m.register_buffer('total_ops', paddle.zeros([1], dtype='int64')) |
| 254 | m.register_buffer('total_params', paddle.zeros([1], dtype='int64')) |
| 255 | m_type = type(m) |
| 256 | |
| 257 | flops_fn = None |
| 258 | if m_type in custom_ops: |
| 259 | flops_fn = custom_ops[m_type] |
| 260 | if m_type not in types_collection: |
| 261 | print(f"Customize Function has been applied to {m_type}") |
| 262 | elif m_type in register_hooks: |
| 263 | flops_fn = register_hooks[m_type] |
| 264 | if m_type not in types_collection: |
| 265 | print(f"{m_type}'s flops has been counted") |
| 266 | else: |
| 267 | if m_type not in types_collection: |
| 268 | print( |
| 269 | f"Cannot find suitable count function for {m_type}. Treat it as zero FLOPs." |
| 270 | ) |
| 271 | |
| 272 | if flops_fn is not None: |
| 273 | flops_handler = m.register_forward_post_hook(flops_fn) |
| 274 | handler_collection.append(flops_handler) |
| 275 | params_handler = m.register_forward_post_hook(count_parameters) |
| 276 | io_handler = m.register_forward_post_hook(count_io_info) |
| 277 | handler_collection.append(params_handler) |
| 278 | handler_collection.append(io_handler) |
| 279 | types_collection.add(m_type) |
| 280 | |
| 281 | training = model.training |
| 282 |
nothing calls this directly
no test coverage detected