(m, x, y)
| 146 | |
| 147 | |
| 148 | def count_convNd(m, x, y): |
| 149 | x = x[0] |
| 150 | kernel_ops = np.prod(m.weight.shape[2:]) |
| 151 | bias_ops = 1 if m.bias is not None else 0 |
| 152 | total_ops = int(y.numel()) * ( |
| 153 | x.shape[1] / m._groups * kernel_ops + bias_ops |
| 154 | ) |
| 155 | m.total_ops += abs(int(total_ops)) |
| 156 | |
| 157 | |
| 158 | def count_leaky_relu(m, x, y): |