@@ -676,6 +676,7 @@ def conv2d(input,
676
676
groups = None ,
677
677
param_attr = None ,
678
678
bias_attr = None ,
679
+ use_cudnn = True ,
679
680
act = None ):
680
681
"""
681
682
**Convlution2D Layer**
@@ -739,6 +740,8 @@ def conv2d(input,
739
740
connected to the second half of the input channels. Default: groups=1
740
741
param_attr(ParamAttr): The parameters to the Conv2d Layer. Default: None
741
742
bias_attr(ParamAttr): Bias parameter for the Conv2d layer. Default: None
743
+ use_cudnn(bool): Use cudnn kernel or not, it is valid only when the cudnn
744
+ library is installed. Default: True
742
745
act(str): Activation type. Default: None
743
746
744
747
Returns:
@@ -774,6 +777,8 @@ def conv2d(input,
774
777
stride = [stride , stride ]
775
778
if isinstance (padding , int ):
776
779
padding = [padding , padding ]
780
+ if not isinstance (use_cudnn , bool ):
781
+ raise ValueError ("use_cudnn should be True or False" )
777
782
778
783
input_shape = input .shape
779
784
filter_shape = [num_filters , num_filter_channels ] + filter_size
@@ -797,9 +802,12 @@ def _get_default_param_initializer():
797
802
'Filter' : filter_param ,
798
803
},
799
804
outputs = {"Output" : pre_bias },
800
- attrs = {'strides' : stride ,
801
- 'paddings' : padding ,
802
- 'groups' : groups })
805
+ attrs = {
806
+ 'strides' : stride ,
807
+ 'paddings' : padding ,
808
+ 'groups' : groups ,
809
+ 'use_cudnn' : use_cudnn
810
+ })
803
811
804
812
pre_act = helper .append_bias_op (pre_bias , dim_start = 1 , dim_end = 2 )
805
813
@@ -948,6 +956,7 @@ def pool2d(input,
948
956
pool_stride = None ,
949
957
pool_padding = None ,
950
958
global_pooling = False ,
959
+ use_cudnn = True ,
951
960
name = None ):
952
961
"""
953
962
This function adds the operator for pooling in 2 dimensions, using the
@@ -967,6 +976,8 @@ def pool2d(input,
967
976
pool_stride = [pool_stride , pool_stride ]
968
977
if isinstance (pool_padding , int ):
969
978
pool_padding = [pool_padding , pool_padding ]
979
+ if not isinstance (use_cudnn , bool ):
980
+ raise ValueError ("use_cudnn should be True or False" )
970
981
971
982
helper = LayerHelper ('pool2d' , ** locals ())
972
983
dtype = helper .input_dtype ()
@@ -981,7 +992,8 @@ def pool2d(input,
981
992
"ksize" : pool_size ,
982
993
"global_pooling" : global_pooling ,
983
994
"strides" : pool_stride ,
984
- "paddings" : pool_padding
995
+ "paddings" : pool_padding ,
996
+ "use_cudnn" : use_cudnn
985
997
})
986
998
987
999
return pool_out
@@ -1096,6 +1108,7 @@ def conv2d_transpose(input,
1096
1108
stride = None ,
1097
1109
dilation = None ,
1098
1110
param_attr = None ,
1111
+ use_cudnn = True ,
1099
1112
name = None ):
1100
1113
"""
1101
1114
The transpose of conv2d layer.
@@ -1123,6 +1136,8 @@ def conv2d_transpose(input,
1123
1136
contain two integers, (dilation_H, dilation_W). Otherwise, the
1124
1137
dilation_H = dilation_W = dilation.
1125
1138
param_attr: Parameter Attribute.
1139
+ use_cudnn(bool): Use cudnn kernel or not, it is valid only when the cudnn
1140
+ library is installed. Default: True
1126
1141
name(str|None): A name for this layer(optional). If set None, the layer
1127
1142
will be named automatically.
1128
1143
@@ -1151,6 +1166,10 @@ def conv2d_transpose(input,
1151
1166
elif dilation is not None :
1152
1167
op_attr ['dilations' ] = dilation
1153
1168
1169
+ if not isinstance (use_cudnn , bool ):
1170
+ raise ValueError ("use_cudnn should be True or False" )
1171
+ op_attr ['use_cudnn' ] = use_cudnn
1172
+
1154
1173
if filter_size is None :
1155
1174
if output_size is None :
1156
1175
raise ValueError ("output_size must be set when filter_size is None" )
0 commit comments