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

1# Copyright 2021 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"""MaxUnpooling2DV2 operation.""" 

16 

17import tensorflow as tf 

18 

19from typeguard import typechecked 

20from typing import Iterable 

21 

22from tensorflow_addons.utils.keras_utils import normalize_tuple 

23 

24 

25def _max_unpooling_2d_v2(updates, mask, output_size): 

26 """Unpool the outputs of a maximum pooling operation.""" 

27 mask = tf.cast(mask, "int32") 

28 input_shape = tf.shape(updates, out_type="int32") 

29 input_shape = [updates.shape[i] or input_shape[i] for i in range(4)] 

30 output_shape = output_size 

31 

32 # Calculates indices for batch, height, width and feature maps. 

33 one_like_mask = tf.ones_like(mask, dtype="int32") 

34 batch_shape = tf.concat([[input_shape[0]], [1], [1], [1]], axis=0) 

35 batch_range = tf.reshape( 

36 tf.range(output_shape[0], dtype="int32"), shape=batch_shape 

37 ) 

38 b = one_like_mask * batch_range 

39 y = mask // (output_shape[2] * output_shape[3]) 

40 x = (mask // output_shape[3]) % output_shape[2] 

41 feature_range = tf.range(output_shape[3], dtype="int32") 

42 f = one_like_mask * feature_range 

43 

44 # Transposes indices & reshape update values to one dimension. 

45 updates_size = tf.size(updates) 

46 indices = tf.transpose(tf.reshape(tf.stack([b, y, x, f]), [4, updates_size])) 

47 values = tf.reshape(updates, [updates_size]) 

48 ret = tf.scatter_nd(indices, values, output_shape) 

49 return ret 

50 

51 

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

53class MaxUnpooling2DV2(tf.keras.layers.Layer): 

54 """Unpool the outputs of a maximum pooling operation. 

55 

56 This differs from MaxUnpooling2D in that it uses output_size rather than strides and padding 

57 to calculate the unpooled tensor. This is because MaxPoolingWithArgMax can map several input 

58 sizes to the same output size, and specifying the output size avoids ambiguity in the 

59 inversion process. 

60 

61 This function currently does not support outputs of MaxPoolingWithArgMax in following cases: 

62 - include_batch_in_index equals true. 

63 - The max pooling operation results in duplicate values in updates and mask. 

64 

65 Args: 

66 output_size: A tuple/list of 4 integers specifying (batch_size, height, width, channel). 

67 The targeted output size. 

68 Call Args: 

69 updates: A 4D tensor of shape `(batch_size, height, width, channel)`. 

70 The pooling result from max pooling. 

71 mask: A 4D tensor of shape `(batch_size, height, width, channel)`. 

72 The indices of the maximal values. 

73 Output shape: 

74 4D tensor with the same shape as output_size. 

75 """ 

76 

77 @typechecked 

78 def __init__( 

79 self, 

80 output_size: Iterable[int], 

81 **kwargs, 

82 ): 

83 super(MaxUnpooling2DV2, self).__init__(**kwargs) 

84 

85 self.output_size = normalize_tuple(output_size, 4, "output_size") 

86 

87 def call(self, updates, mask): 

88 return _max_unpooling_2d_v2(updates, mask, output_size=self.output_size) 

89 

90 def get_config(self): 

91 config = super(MaxUnpooling2DV2, self).get_config() 

92 config["output_size"] = self.output_size 

93 return config