Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/tensorflow_addons/layers/adaptive_pooling.py: 36%

1# Copyright 2020 The TensorFlow Authors. All Rights Reserved. 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================== 

15"""Pooling layers with fixed size outputs""" 

16 

17import tensorflow as tf 

18import tensorflow_addons.utils.keras_utils as conv_utils 

19 

20from typeguard import typechecked 

21from typing import Union, Callable, Iterable 

22 

23 

24class AdaptivePooling1D(tf.keras.layers.Layer): 

25 """Parent class for 1D pooling layers with adaptive kernel size. 

26 

27 This class only exists for code reuse. It will never be an exposed API. 

28 

29 Args: 

30 reduce_function: The reduction method to apply, e.g. `tf.reduce_max`. 

31 output_size: An integer or tuple/list of a single integer, specifying pooled_features. 

32 The new size of output channels. 

33 data_format: A string, 

34 one of `channels_last` (default) or `channels_first`. 

35 The ordering of the dimensions in the inputs. 

36 `channels_last` corresponds to inputs with shape 

37 `(batch, steps, features)` while `channels_first` 

38 corresponds to inputs with shape 

39 `(batch, features, steps)`. 

40 """ 

41 

42 @typechecked 

43 def __init__( 

44 self, 

45 reduce_function: Callable, 

46 output_size: Union[int, Iterable[int]], 

47 data_format=None, 

48 **kwargs, 

49 ): 

50 self.data_format = conv_utils.normalize_data_format(data_format) 

51 self.reduce_function = reduce_function 

52 self.output_size = conv_utils.normalize_tuple(output_size, 1, "output_size") 

53 super().__init__(**kwargs) 

54 

55 def call(self, inputs, *args): 

56 bins = self.output_size[0] 

57 if self.data_format == "channels_last": 

58 splits = tf.split(inputs, bins, axis=1) 

59 splits = tf.stack(splits, axis=1) 

60 out_vect = self.reduce_function(splits, axis=2) 

61 else: 

62 splits = tf.split(inputs, bins, axis=2) 

63 splits = tf.stack(splits, axis=2) 

64 out_vect = self.reduce_function(splits, axis=3) 

65 return out_vect 

66 

67 def compute_output_shape(self, input_shape): 

68 input_shape = tf.TensorShape(input_shape).as_list() 

69 if self.data_format == "channels_last": 

70 shape = tf.TensorShape( 

71 [input_shape[0], self.output_size[0], input_shape[2]] 

72 ) 

73 else: 

74 shape = tf.TensorShape( 

75 [input_shape[0], input_shape[1], self.output_size[0]] 

76 ) 

77 

78 return shape 

79 

80 def get_config(self): 

81 config = { 

82 "output_size": self.output_size, 

83 "data_format": self.data_format, 

84 } 

85 base_config = super().get_config() 

86 return {**base_config, **config} 

87 

88 

89@tf.keras.utils.register_keras_serializable(package="Addons") 

90class AdaptiveAveragePooling1D(AdaptivePooling1D): 

91 """Average Pooling with adaptive kernel size. 

92 

93 Args: 

94 output_size: An integer or tuple/list of a single integer, specifying pooled_features. 

95 The new size of output channels. 

96 data_format: A string, 

97 one of `channels_last` (default) or `channels_first`. 

98 The ordering of the dimensions in the inputs. 

99 `channels_last` corresponds to inputs with shape 

100 `(batch, steps, channels)` while `channels_first` 

101 corresponds to inputs with shape `(batch, channels, steps)`. 

102 

103 Input shape: 

104 - If `data_format='channels_last'`: 

105 3D tensor with shape `(batch, steps, channels)`. 

106 - If `data_format='channels_first'`: 

107 3D tensor with shape `(batch, channels, steps)`. 

108 

109 Output shape: 

110 - If `data_format='channels_last'`: 

111 3D tensor with shape `(batch_size, pooled_steps, channels)`. 

112 - If `data_format='channels_first'`: 

113 3D tensor with shape `(batch_size, channels, pooled_steps)`. 

114 """ 

115 

116 @typechecked 

117 def __init__( 

118 self, output_size: Union[int, Iterable[int]], data_format=None, **kwargs 

119 ): 

120 super().__init__(tf.reduce_mean, output_size, data_format, **kwargs) 

121 

122 

123@tf.keras.utils.register_keras_serializable(package="Addons") 

124class AdaptiveMaxPooling1D(AdaptivePooling1D): 

125 """Max Pooling with adaptive kernel size. 

126 

127 Args: 

128 output_size: An integer or tuple/list of a single integer, specifying pooled_features. 

129 The new size of output channels. 

130 data_format: A string, 

131 one of `channels_last` (default) or `channels_first`. 

132 The ordering of the dimensions in the inputs. 

133 `channels_last` corresponds to inputs with shape 

134 `(batch, steps, channels)` while `channels_first` 

135 corresponds to inputs with shape `(batch, channels, steps)`. 

136 

137 Input shape: 

138 - If `data_format='channels_last'`: 

139 3D tensor with shape `(batch, steps, channels)`. 

140 - If `data_format='channels_first'`: 

141 3D tensor with shape `(batch, channels, steps)`. 

142 

143 Output shape: 

144 - If `data_format='channels_last'`: 

145 3D tensor with shape `(batch_size, pooled_steps, channels)`. 

146 - If `data_format='channels_first'`: 

147 3D tensor with shape `(batch_size, channels, pooled_steps)`. 

148 """ 

149 

150 @typechecked 

151 def __init__( 

152 self, output_size: Union[int, Iterable[int]], data_format=None, **kwargs 

153 ): 

154 super().__init__(tf.reduce_max, output_size, data_format, **kwargs) 

155 

156 

157class AdaptivePooling2D(tf.keras.layers.Layer): 

158 """Parent class for 2D pooling layers with adaptive kernel size. 

159 

160 This class only exists for code reuse. It will never be an exposed API. 

161 

162 Args: 

163 reduce_function: The reduction method to apply, e.g. `tf.reduce_max`. 

164 output_size: An integer or tuple/list of 2 integers specifying (pooled_rows, pooled_cols). 

165 The new size of output channels. 

166 data_format: A string, 

167 one of `channels_last` (default) or `channels_first`. 

168 The ordering of the dimensions in the inputs. 

169 `channels_last` corresponds to inputs with shape 

170 `(batch, height, width, channels)` while `channels_first` 

171 corresponds to inputs with shape 

172 `(batch, channels, height, width)`. 

173 """ 

174 

175 @typechecked 

176 def __init__( 

177 self, 

178 reduce_function: Callable, 

179 output_size: Union[int, Iterable[int]], 

180 data_format=None, 

181 **kwargs, 

182 ): 

183 self.data_format = conv_utils.normalize_data_format(data_format) 

184 self.reduce_function = reduce_function 

185 self.output_size = conv_utils.normalize_tuple(output_size, 2, "output_size") 

186 super().__init__(**kwargs) 

187 

188 def call(self, inputs, *args): 

189 h_bins = self.output_size[0] 

190 w_bins = self.output_size[1] 

191 if self.data_format == "channels_last": 

192 split_cols = tf.split(inputs, h_bins, axis=1) 

193 split_cols = tf.stack(split_cols, axis=1) 

194 split_rows = tf.split(split_cols, w_bins, axis=3) 

195 split_rows = tf.stack(split_rows, axis=3) 

196 out_vect = self.reduce_function(split_rows, axis=[2, 4]) 

197 else: 

198 split_cols = tf.split(inputs, h_bins, axis=2) 

199 split_cols = tf.stack(split_cols, axis=2) 

200 split_rows = tf.split(split_cols, w_bins, axis=4) 

201 split_rows = tf.stack(split_rows, axis=4) 

202 out_vect = self.reduce_function(split_rows, axis=[3, 5]) 

203 return out_vect 

204 

205 def compute_output_shape(self, input_shape): 

206 input_shape = tf.TensorShape(input_shape).as_list() 

207 if self.data_format == "channels_last": 

208 shape = tf.TensorShape( 

209 [ 

210 input_shape[0], 

211 self.output_size[0], 

212 self.output_size[1], 

213 input_shape[3], 

214 ] 

215 ) 

216 else: 

217 shape = tf.TensorShape( 

218 [ 

219 input_shape[0], 

220 input_shape[1], 

221 self.output_size[0], 

222 self.output_size[1], 

223 ] 

224 ) 

225 

226 return shape 

227 

228 def get_config(self): 

229 config = { 

230 "output_size": self.output_size, 

231 "data_format": self.data_format, 

232 } 

233 base_config = super().get_config() 

234 return {**base_config, **config} 

235 

236 

237@tf.keras.utils.register_keras_serializable(package="Addons") 

238class AdaptiveAveragePooling2D(AdaptivePooling2D): 

239 """Average Pooling with adaptive kernel size. 

240 

241 Args: 

242 output_size: Tuple of integers specifying (pooled_rows, pooled_cols). 

243 The new size of output channels. 

244 data_format: A string, 

245 one of `channels_last` (default) or `channels_first`. 

246 The ordering of the dimensions in the inputs. 

247 `channels_last` corresponds to inputs with shape 

248 `(batch, height, width, channels)` while `channels_first` 

249 corresponds to inputs with shape `(batch, channels, height, width)`. 

250 

251 Input shape: 

252 - If `data_format='channels_last'`: 

253 4D tensor with shape `(batch_size, height, width, channels)`. 

254 - If `data_format='channels_first'`: 

255 4D tensor with shape `(batch_size, channels, height, width)`. 

256 

257 Output shape: 

258 - If `data_format='channels_last'`: 

259 4D tensor with shape `(batch_size, pooled_rows, pooled_cols, channels)`. 

260 - If `data_format='channels_first'`: 

261 4D tensor with shape `(batch_size, channels, pooled_rows, pooled_cols)`. 

262 """ 

263 

264 @typechecked 

265 def __init__( 

266 self, output_size: Union[int, Iterable[int]], data_format=None, **kwargs 

267 ): 

268 super().__init__(tf.reduce_mean, output_size, data_format, **kwargs) 

269 

270 

271@tf.keras.utils.register_keras_serializable(package="Addons") 

272class AdaptiveMaxPooling2D(AdaptivePooling2D): 

273 """Max Pooling with adaptive kernel size. 

274 

275 Args: 

276 output_size: Tuple of integers specifying (pooled_rows, pooled_cols). 

277 The new size of output channels. 

278 data_format: A string, 

279 one of `channels_last` (default) or `channels_first`. 

280 The ordering of the dimensions in the inputs. 

281 `channels_last` corresponds to inputs with shape 

282 `(batch, height, width, channels)` while `channels_first` 

283 corresponds to inputs with shape `(batch, channels, height, width)`. 

284 

285 Input shape: 

286 - If `data_format='channels_last'`: 

287 4D tensor with shape `(batch_size, height, width, channels)`. 

288 - If `data_format='channels_first'`: 

289 4D tensor with shape `(batch_size, channels, height, width)`. 

290 

291 Output shape: 

292 - If `data_format='channels_last'`: 

293 4D tensor with shape `(batch_size, pooled_rows, pooled_cols, channels)`. 

294 - If `data_format='channels_first'`: 

295 4D tensor with shape `(batch_size, channels, pooled_rows, pooled_cols)`. 

296 """ 

297 

298 @typechecked 

299 def __init__( 

300 self, output_size: Union[int, Iterable[int]], data_format=None, **kwargs 

301 ): 

302 super().__init__(tf.reduce_max, output_size, data_format, **kwargs) 

303 

304 

305class AdaptivePooling3D(tf.keras.layers.Layer): 

306 """Parent class for 3D pooling layers with adaptive kernel size. 

307 

308 This class only exists for code reuse. It will never be an exposed API. 

309 

310 Args: 

311 reduce_function: The reduction method to apply, e.g. `tf.reduce_max`. 

312 output_size: An integer or tuple/list of 3 integers specifying (pooled_dim1, pooled_dim2, pooled_dim3). 

313 The new size of output channels. 

314 data_format: A string, 

315 one of `channels_last` (default) or `channels_first`. 

316 The ordering of the dimensions in the inputs. 

317 `channels_last` corresponds to inputs with shape 

318 `(batch, spatial_dim1, spatial_dim2, spatial_dim3, channels)` while `channels_first` 

319 corresponds to inputs with shape 

320 `(batch, channels, spatial_dim1, spatial_dim2, spatial_dim3)`. 

321 """ 

322 

323 @typechecked 

324 def __init__( 

325 self, 

326 reduce_function: Callable, 

327 output_size: Union[int, Iterable[int]], 

328 data_format=None, 

329 **kwargs, 

330 ): 

331 self.data_format = conv_utils.normalize_data_format(data_format) 

332 self.reduce_function = reduce_function 

333 self.output_size = conv_utils.normalize_tuple(output_size, 3, "output_size") 

334 super().__init__(**kwargs) 

335 

336 def call(self, inputs, *args): 

337 h_bins = self.output_size[0] 

338 w_bins = self.output_size[1] 

339 d_bins = self.output_size[2] 

340 if self.data_format == "channels_last": 

341 split_cols = tf.split(inputs, h_bins, axis=1) 

342 split_cols = tf.stack(split_cols, axis=1) 

343 split_rows = tf.split(split_cols, w_bins, axis=3) 

344 split_rows = tf.stack(split_rows, axis=3) 

345 split_depth = tf.split(split_rows, d_bins, axis=5) 

346 split_depth = tf.stack(split_depth, axis=5) 

347 out_vect = self.reduce_function(split_depth, axis=[2, 4, 6]) 

348 else: 

349 split_cols = tf.split(inputs, h_bins, axis=2) 

350 split_cols = tf.stack(split_cols, axis=2) 

351 split_rows = tf.split(split_cols, w_bins, axis=4) 

352 split_rows = tf.stack(split_rows, axis=4) 

353 split_depth = tf.split(split_rows, d_bins, axis=6) 

354 split_depth = tf.stack(split_depth, axis=6) 

355 out_vect = self.reduce_function(split_depth, axis=[3, 5, 7]) 

356 return out_vect 

357 

358 def compute_output_shape(self, input_shape): 

359 input_shape = tf.TensorShape(input_shape).as_list() 

360 if self.data_format == "channels_last": 

361 shape = tf.TensorShape( 

362 [ 

363 input_shape[0], 

364 self.output_size[0], 

365 self.output_size[1], 

366 self.output_size[2], 

367 input_shape[4], 

368 ] 

369 ) 

370 else: 

371 shape = tf.TensorShape( 

372 [ 

373 input_shape[0], 

374 input_shape[1], 

375 self.output_size[0], 

376 self.output_size[1], 

377 self.output_size[2], 

378 ] 

379 ) 

380 

381 return shape 

382 

383 def get_config(self): 

384 config = { 

385 "output_size": self.output_size, 

386 "data_format": self.data_format, 

387 } 

388 base_config = super().get_config() 

389 return {**base_config, **config} 

390 

391 

392@tf.keras.utils.register_keras_serializable(package="Addons") 

393class AdaptiveAveragePooling3D(AdaptivePooling3D): 

394 """Average Pooling with adaptive kernel size. 

395 

396 Args: 

397 output_size: An integer or tuple/list of 3 integers specifying (pooled_depth, pooled_height, pooled_width). 

398 The new size of output channels. 

399 data_format: A string, 

400 one of `channels_last` (default) or `channels_first`. 

401 The ordering of the dimensions in the inputs. 

402 `channels_last` corresponds to inputs with shape 

403 `(batch, height, width, channels)` while `channels_first` 

404 corresponds to inputs with shape `(batch, channels, height, width)`. 

405 

406 Input shape: 

407 - If `data_format='channels_last'`: 

408 5D tensor with shape `(batch_size, spatial_dim1, spatial_dim2, spatial_dim3, channels)`. 

409 - If `data_format='channels_first'`: 

410 5D tensor with shape `(batch_size, channels, spatial_dim1, spatial_dim2, spatial_dim3)`. 

411 

412 Output shape: 

413 - If `data_format='channels_last'`: 

414 5D tensor with shape `(batch_size, pooled_dim1, pooled_dim2, pooled_dim3, channels)`. 

415 - If `data_format='channels_first'`: 

416 5D tensor with shape `(batch_size, channels, pooled_dim1, pooled_dim2, pooled_dim3)`. 

417 """ 

418 

419 @typechecked 

420 def __init__( 

421 self, output_size: Union[int, Iterable[int]], data_format=None, **kwargs 

422 ): 

423 super().__init__(tf.reduce_mean, output_size, data_format, **kwargs) 

424 

425 

426@tf.keras.utils.register_keras_serializable(package="Addons") 

427class AdaptiveMaxPooling3D(AdaptivePooling3D): 

428 """Max Pooling with adaptive kernel size. 

429 

430 Args: 

431 output_size: An integer or tuple/list of 3 integers specifying (pooled_depth, pooled_height, pooled_width). 

432 The new size of output channels. 

433 data_format: A string, 

434 one of `channels_last` (default) or `channels_first`. 

435 The ordering of the dimensions in the inputs. 

436 `channels_last` corresponds to inputs with shape 

437 `(batch, height, width, channels)` while `channels_first` 

438 corresponds to inputs with shape `(batch, channels, height, width)`. 

439 

440 Input shape: 

441 - If `data_format='channels_last'`: 

442 5D tensor with shape `(batch_size, spatial_dim1, spatial_dim2, spatial_dim3, channels)`. 

443 - If `data_format='channels_first'`: 

444 5D tensor with shape `(batch_size, channels, spatial_dim1, spatial_dim2, spatial_dim3)`. 

445 

446 Output shape: 

447 - If `data_format='channels_last'`: 

448 5D tensor with shape `(batch_size, pooled_dim1, pooled_dim2, pooled_dim3, channels)`. 

449 - If `data_format='channels_first'`: 

450 5D tensor with shape `(batch_size, channels, pooled_dim1, pooled_dim2, pooled_dim3)`. 

451 """ 

452 

453 @typechecked 

454 def __init__( 

455 self, output_size: Union[int, Iterable[int]], data_format=None, **kwargs 

456 ): 

457 super().__init__(tf.reduce_max, output_size, data_format, **kwargs)