A counter for the `thop` package to count the operations in an attention operation. Meant to be used like: macs, params = thop.profile( model, inputs=(inputs, timestamps), custom_ops={QKVAttention: QKVAttention.count_flops}, )
(model, _x, y)
| 326 | |
| 327 | |
| 328 | def count_flops_attn(model, _x, y): |
| 329 | """ |
| 330 | A counter for the `thop` package to count the operations in an |
| 331 | attention operation. |
| 332 | Meant to be used like: |
| 333 | macs, params = thop.profile( |
| 334 | model, |
| 335 | inputs=(inputs, timestamps), |
| 336 | custom_ops={QKVAttention: QKVAttention.count_flops}, |
| 337 | ) |
| 338 | """ |
| 339 | b, c, *spatial = y[0].shape |
| 340 | num_spatial = int(np.prod(spatial)) |
| 341 | # We perform two matmuls with the same number of ops. |
| 342 | # The first computes the weight matrix, the second computes |
| 343 | # the combination of the value vectors. |
| 344 | matmul_ops = 2 * b * (num_spatial ** 2) * c |
| 345 | model.total_ops += th.DoubleTensor([matmul_ops]) |
| 346 | |
| 347 | |
| 348 | class QKVAttentionLegacy(nn.Module): |
no outgoing calls
no test coverage detected