[fp16] suppot fp16 on nn.Dropout2D (#50904)

* add unittest for nn.DropOut2D

* add fp16

* add fp16 in docs of temporal_shift_op.cc

* Update test_dropout_op.py

paddle/fluid/operators/temporal_shift_op.cc

@@ -45,7 +45,7 @@ class TemporalShiftOpMaker : public framework::OpProtoAndCheckerMaker {
              "While N is the batch size, T is the temporal segment "
              "number, C is the channel number, H is the height of "
              "features and W is the width of features. "
-             "The data type is float32 and float64");
+             "The data type is float16, float32 and float64");
     AddOutput("Out",
               "The output tensor of temporal shift operator. "
               "This is a 4-D tensor in the same shape with Input(X).");

python/paddle/fluid/tests/unittests/test_dropout_op.py

@@ -769,6 +769,29 @@ class TestDropout2DCAPI(unittest.TestCase):
                     result.numpy(), result_np, rtol=1e-05
                 )
 
+    def test_static_fp16_with_gpu(self):
+        if paddle.fluid.core.is_compiled_with_cuda():
+            place = paddle.CUDAPlace(0)
+            with paddle.static.program_guard(
+                paddle.static.Program(), paddle.static.Program()
+            ):
+                input = paddle.static.data(
+                    name="input", shape=[2, 3, 4, 5], dtype="float16"
+                )
+
+                m = paddle.nn.Dropout2D(p=0.5)
+                res1 = m(input)
+
+                in_np = np.random.random([2, 3, 4, 5]).astype("float16")
+                res_np = in_np
+
+                exe = paddle.static.Executor(place)
+                fetches = exe.run(
+                    paddle.static.default_main_program(),
+                    feed={"input": in_np},
+                    fetch_list=[res1],
+                )
+
 
 class TestDropout3DFAPI(unittest.TestCase):
     def setUp(self):

python/paddle/nn/functional/common.py

@@ -982,7 +982,7 @@ def dropout(
     dropout probability.
 
     Args:
-        x (Tensor): The input tensor. The data type is float32 or float64.
+        x (Tensor): The input tensor. The data type is float16, float32 or float64.
         p (float|int, optional): Probability of setting units to zero. Default: 0.5.
         axis (int|list|tuple, optional): The axis along which the dropout is performed. Default: None.
         training (bool, optional): A flag indicating whether it is in train phrase or not. Default: True.
@@ -1201,7 +1201,9 @@ def dropout(
             return out
     else:  # sometimes called dropout_nd #TODO: optimize with c++
         if not in_dynamic_mode():
-            check_variable_and_dtype(x, 'x', ['float32', 'float64'], 'dropout')
+            check_variable_and_dtype(
+                x, 'x', ['float16', 'float32', 'float64'], 'dropout'
+            )
         dtype = x.dtype
         keep_prob = 1 - p
         if training:
@@ -1269,7 +1271,7 @@ def dropout2d(x, p=0.5, training=True, data_format='NCHW', name=None):
 
     Args:
         x (Tensor):  The input is 4-D Tensor with shape [N, C, H, W] or [N, H, W, C].
-                     The data type is float32 or float64.
+                     The data type is float16, float32 or float64.
         p (float, optional): Probability of setting units to zero. Default: 0.5.
         training (bool, optional): A flag indicating whether it is in train phrase or not. Default: True.
         data_format (str, optional): Specify the data format of the input, and the data format of the output will be consistent with that of the input. An optional string from `NCHW` or `NHWC` . When it is `NCHW` , the data is stored in the order of: [batch_size, input_channels, input_height, input_width]. Default: `NCHW` .