Commit c575973
authored
[Relax][Onnx][BatchNorm] Pass momentum and training_mode into BatchNorm Operator (#18704)
### Description
- Onnx model have training_mode atrr = False, but Relax model after
convert have training = True
- Momentum values in Relax module are not the same as onnx model
### Steps to Reproduce
<img width="600" height="400" alt="BatchNorm"
src="https://github.com/user-attachments/assets/2f0ca26b-e83b-4ab8-ab06-a537802af6de"
/>
- Relax model:
```
class Module:
def main(X: R.Tensor((2, 3, 4, 4), dtype="float32")) -> R.Tensor((2, 3, 4, 4), dtype="float32"):
R.func_attr({"num_input": 1})
with R.dataflow():
lv: R.Tuple(R.Tensor((2, 3, 4, 4), dtype="float32"), R.Tensor((3,), dtype="float32"), R.Tensor((3,), dtype="float32")) = R.nn.batch_norm(X, metadata["relax.expr.Constant"][0], metadata["relax.expr.Constant"][1], metadata["relax.expr.Constant"][2], metadata["relax.expr.Constant"][3], axis=1, epsilon=9.9999997473787516e-06, center=True, scale=True, momentum=0.10000000000000001, training=True)
lv1: R.Tensor((2, 3, 4, 4), dtype="float32") = lv[0]
lv2: R.Tensor((3,), dtype="float32") = lv[1]
lv3: R.Tensor((3,), dtype="float32") = lv[2]
gv: R.Tensor((2, 3, 4, 4), dtype="float32") = lv1
R.output(gv)
return gv
```
### Resolved
- Get Attributes and Pass momentum/training_mode with default value into
BatchNorm Operator
- Fixed: #187031 parent 4b2b639 commit c575973
1 file changed
Lines changed: 11 additions & 1 deletion
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
2435 | 2435 | | |
2436 | 2436 | | |
2437 | 2437 | | |
| 2438 | + | |
| 2439 | + | |
2438 | 2440 | | |
2439 | | - | |
| 2441 | + | |
| 2442 | + | |
| 2443 | + | |
| 2444 | + | |
| 2445 | + | |
| 2446 | + | |
| 2447 | + | |
| 2448 | + | |
| 2449 | + | |
2440 | 2450 | | |
2441 | 2451 | | |
2442 | 2452 | | |
| |||
0 commit comments