add module.zero_grad

[Flowingsun007]

No description provided.

[daquexian]

添加一个测试用例吧

[Flowingsun007]

添加一个测试用例吧

好的，我后面补一个

[github-actions]

Code got formatted by CI. Please request CI again if you still want to have this PR merged. If the PR is from a forked repo, please download the patch files from the GitHub Actions web page and apply them locally.

[github-actions]

Speed stats:

GPU Name: GeForce GTX 1080 

✔️ OneFlow resnet50 time: 128.7ms (= 12868.1ms / 100, input_shape=[16, 3, 224, 224])
PyTorch resnet50 time: 137.5ms (= 13745.9ms / 100, input_shape=[16, 3, 224, 224])
✔️ Relative speed: 1.07 (= 137.5ms / 128.7ms)

✔️ OneFlow resnet50 time: 76.9ms (= 7693.7ms / 100, input_shape=[8, 3, 224, 224])
PyTorch resnet50 time: 84.8ms (= 8478.2ms / 100, input_shape=[8, 3, 224, 224])
✔️ Relative speed: 1.10 (= 84.8ms / 76.9ms)

OneFlow resnet50 time: 53.2ms (= 10635.3ms / 200, input_shape=[4, 3, 224, 224])
PyTorch resnet50 time: 56.4ms (= 11271.2ms / 200, input_shape=[4, 3, 224, 224])
✔️ Relative speed: 1.06 (= 56.4ms / 53.2ms)

OneFlow resnet50 time: 44.2ms (= 8847.0ms / 200, input_shape=[2, 3, 224, 224])
PyTorch resnet50 time: 49.1ms (= 9828.4ms / 200, input_shape=[2, 3, 224, 224])
✔️ Relative speed: 1.11 (= 49.1ms / 44.2ms)

OneFlow resnet50 time: 38.1ms (= 7610.5ms / 200, input_shape=[1, 3, 224, 224])
PyTorch resnet50 time: 38.9ms (= 7776.2ms / 200, input_shape=[1, 3, 224, 224])
✔️ Relative speed: 1.02 (= 38.9ms / 38.1ms)

✔️ OneFlow resnet50 time: 142.8ms (= 14276.3ms / 100, input_shape=[16, 3, 224, 224], ddp, world size=2)
PyTorch resnet50 time: 162.6ms (= 16260.6ms / 100, input_shape=[16, 3, 224, 224], ddp, world size=2)
✔️ Relative speed: 1.14 (= 162.6ms / 142.8ms)

OneFlow resnet50 time: 90.2ms (= 9020.8ms / 100, input_shape=[8, 3, 224, 224], ddp, world size=2)
PyTorch resnet50 time: 101.4ms (= 10138.3ms / 100, input_shape=[8, 3, 224, 224], ddp, world size=2)
✔️ Relative speed: 1.12 (= 101.4ms / 90.2ms)

OneFlow resnet50 time: 61.9ms (= 12382.8ms / 200, input_shape=[4, 3, 224, 224], ddp, world size=2)
PyTorch resnet50 time: 75.8ms (= 15169.0ms / 200, input_shape=[4, 3, 224, 224], ddp, world size=2)
✔️ Relative speed: 1.23 (= 75.8ms / 61.9ms)

OneFlow resnet50 time: 53.8ms (= 10753.8ms / 200, input_shape=[2, 3, 224, 224], ddp, world size=2)
PyTorch resnet50 time: 66.4ms (= 13277.3ms / 200, input_shape=[2, 3, 224, 224], ddp, world size=2)
✔️ Relative speed: 1.23 (= 66.4ms / 53.8ms)

OneFlow resnet50 time: 46.6ms (= 9326.5ms / 200, input_shape=[1, 3, 224, 224], ddp, world size=2)
PyTorch resnet50 time: 58.0ms (= 11598.8ms / 200, input_shape=[1, 3, 224, 224], ddp, world size=2)
✔️ Relative speed: 1.24 (= 58.0ms / 46.6ms)