wrap flow.nn.init.* with flow.no_grad() (#5932)

* wrap flow.nn.init.* with flow.no_grad()
Signed-off-by: Ndaquexian <daquexian566@gmail.com>

* auto format by CI
Co-authored-by: Noneflow-ci-bot <ci-bot@oneflow.org>
Co-authored-by: Noneflow-ci-bot <69100618+oneflow-ci-bot@users.noreply.github.com>

python/oneflow/framework/tensor.py

@@ -242,7 +242,7 @@ def _pow(self, b):
     return flow.pow(self, b)
 
 
-def _uniform_(self, a=0, b=1):
+def _uniform(self, a=0, b=1):
     initializer_conf = flow.random_uniform_initializer(
         minval=a, maxval=b, dtype=self.dtype
     )
@@ -396,7 +396,7 @@ def RegisterMethods():
     Tensor.__neg__ = _neg
     Tensor.__pow__ = _pow
     Tensor.__format__ = _format
-    Tensor.uniform_ = _uniform_
+    Tensor.uniform_ = _uniform
     Tensor.kaiming_uniform_ = _kaiming_uniform
     Tensor.kaiming_normal_ = _kaiming_normal
     Tensor.xavier_normal_ = _xavier_normal

python/oneflow/nn/init.py

@@ -13,6 +13,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+import oneflow as flow
 from oneflow.ops.initializer_util import CalcGain
 
 
@@ -21,43 +22,52 @@ def calculate_gain(nonlinearity, param=None):
 
 
 def uniform_(tensor, a=0.0, b=1.0):
-    tensor.uniform_(a, b)
+    with flow.no_grad():
+        return tensor.uniform_(a, b)
 
 
 def normal_(tensor, mean=0.0, std=1.0):
-    tensor.normal_(mean, std)
+    with flow.no_grad():
+        return tensor.normal_(mean, std)
 
 
 def xavier_uniform_(tensor, gain=1.0, *, data_format="NCHW"):
-    tensor.xavier_uniform_(gain, data_format=data_format)
+    with flow.no_grad():
+        return tensor.xavier_uniform_(gain, data_format=data_format)
 
 
 def xavier_normal_(tensor, gain=1.0, *, data_format="NCHW"):
-    tensor.xavier_normal_(gain, data_format=data_format)
+    with flow.no_grad():
+        return tensor.xavier_normal_(gain, data_format=data_format)
 
 
 def kaiming_uniform_(
     tensor, a=0, mode="fan_in", nonlinearity="leaky_relu", *, data_format="NCHW"
 ):
-    tensor.kaiming_uniform_(a, mode, nonlinearity, data_format=data_format)
+    with flow.no_grad():
+        return tensor.kaiming_uniform_(a, mode, nonlinearity, data_format=data_format)
 
 
 def kaiming_normal_(
     tensor, a=0, mode="fan_in", nonlinearity="leaky_relu", *, data_format="NCHW"
 ):
-    tensor.kaiming_normal_(a, mode, nonlinearity, data_format=data_format)
+    with flow.no_grad():
+        return tensor.kaiming_normal_(a, mode, nonlinearity, data_format=data_format)
 
 
 def constant_(tensor, val):
-    tensor.fill_(val)
+    with flow.no_grad():
+        return tensor.fill_(val)
 
 
 def ones_(tensor):
-    tensor.fill_(1)
+    with flow.no_grad():
+        return tensor.fill_(1)
 
 
 def zeros_(tensor):
-    tensor.fill_(0)
+    with flow.no_grad():
+        return tensor.fill_(0)
 
 
 def _calculate_fan_in_and_fan_out(tensor):

python/oneflow/test/tensor/test_tensor.py

@@ -99,6 +99,7 @@ class TestTensor(flow.unittest.TestCase):
         z = tensor_creator(5, 4, 3, 2)
         flow.nn.init.kaiming_normal_(z, a=0.1, mode="fan_out", nonlinearity="relu")
         flow.nn.init.kaiming_uniform_(z)
+        z.requires_grad_()
         flow.nn.init.xavier_normal_(z)
         flow.nn.init.xavier_uniform_(z)
         x = tensor_creator(*shape, dtype=flow.int32)