Get output shape of 2D or 3D convolution Paramters --------- tensor_format: int 0: CUDNN_TENSOR_NCHW 1: CUDNN_TENSOR_NHWC 2: CUDNN_TENSOR_NCHW_VECT_C pad: int or list padding stride: int or list stride dilation: int or list dil
(
tensor_format, pad, stride, dilation, x_shape, w_shape, data_dtype, conv_dtype, groups=1
)
| 141 | |
| 142 | |
| 143 | def conv_output_shape( |
| 144 | tensor_format, pad, stride, dilation, x_shape, w_shape, data_dtype, conv_dtype, groups=1 |
| 145 | ): |
| 146 | """Get output shape of 2D or 3D convolution |
| 147 | |
| 148 | Paramters |
| 149 | --------- |
| 150 | tensor_format: int |
| 151 | 0: CUDNN_TENSOR_NCHW |
| 152 | 1: CUDNN_TENSOR_NHWC |
| 153 | 2: CUDNN_TENSOR_NCHW_VECT_C |
| 154 | pad: int or list |
| 155 | padding |
| 156 | stride: int or list |
| 157 | stride |
| 158 | dilation: int or list |
| 159 | dilation |
| 160 | x_shape: list |
| 161 | input shape |
| 162 | w_shape: list |
| 163 | weight shape |
| 164 | data_dtype: str |
| 165 | data type |
| 166 | conv_dtype: str |
| 167 | convolution type |
| 168 | groups: int |
| 169 | number of groups |
| 170 | |
| 171 | Returns |
| 172 | ------- |
| 173 | oshape: list |
| 174 | output shape |
| 175 | """ |
| 176 | |
| 177 | assert len(x_shape) == len(w_shape) |
| 178 | assert len(x_shape) in (4, 5) |
| 179 | |
| 180 | if tensor_format == 0: |
| 181 | n_output = x_shape[0] |
| 182 | c_output = w_shape[0] |
| 183 | x_chan = x_shape[1] |
| 184 | w_chan_input = w_shape[1] |
| 185 | x_shape = x_shape[2:] |
| 186 | w_shape = w_shape[2:] |
| 187 | |
| 188 | elif tensor_format == 1: |
| 189 | n_output = x_shape[0] |
| 190 | c_output = w_shape[0] |
| 191 | x_chan = x_shape[-1] |
| 192 | w_chan_input = w_shape[-1] |
| 193 | assert len(x_shape) == 4, "CuDNN layout NHWC is only well-defined for 4d tensors" |
| 194 | x_shape = x_shape[1:-1] |
| 195 | w_shape = w_shape[1:-1] |
| 196 | |
| 197 | elif tensor_format == 2: |
| 198 | n_output = x_shape[0] |
| 199 | c_output = w_shape[0] |
| 200 | x_chan = x_shape[1] |
no test coverage detected
searching dependent graphs…