tensorflow改写torch的AdaptiveAvgPool2d

tensorflow改写torch的AdaptiveAvgPool2d代码实例

import tensorflow as tf
import torch as t
import numpy as np
tf.keras.backend.set_floatx('float32')
input_data = np.random.rand(1, 7, 7, 2)  # (B, H, W, C) 这是temsorflow中的顺序
# pytorch tensor的通道排序 [batch, channel, height, width]
pytorch_input = t.tensor(input_data, dtype=t.float32).permute(0, 3, 1, 2)
print(pytorch_input.shape)
tensorflow_input = tf.convert_to_tensor(input_data)
print(tensorflow_input.shape)

inputsz = 7
outputsz = 2
stridesz = np.floor(inputsz / outputsz).astype(np.int32)  # 向下取整
print("stridesz:", stridesz)
kernelsz = inputsz - (outputsz - 1) * stridesz


adp = t.nn.AdaptiveAvgPool2d(output_size=outputsz)
avg = tf.keras.layers.AveragePooling2D((kernelsz, kernelsz), strides=(stridesz, stridesz))

adplist = adp(pytorch_input)
avglist = avg(tensorflow_input)
print(input_data)
print(adplist)
print(adplist.shape) # (B, C, H, w)
## 把torch的顺序改的和tensorflow一致
print(adplist.permute(0, 2, 3, 1)) # (B, H, W, C)
print(avglist)

D:\Apps\Anaconda3\envs\tensorflow-env\python.exe D:/PYproject/start/algorithmlib/Test/test4.py
torch.Size([1, 2, 7, 7])
(1, 7, 7, 2)
stridesz: 3
2022-07-19 11:28:15.407475: I tensorflow/core/platform/cpu_feature_guard.cc:142] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX AVX2
2022-07-19 11:28:15.410292: I tensorflow/core/common_runtime/process_util.cc:147] Creating new thread pool with default inter op setting: 2. Tune using inter_op_parallelism_threads for best performance.
WARNING:tensorflow:Layer average_pooling2d is casting an input tensor from dtype float64 to the layer's dtype of float32, which is new behavior in TensorFlow 2.  The layer has dtype float32 because it's dtype defaults to floatx.

If you intended to run this layer in float32, you can safely ignore this warning. If in doubt, this warning is likely only an issue if you are porting a TensorFlow 1.X model to TensorFlow 2.

To change all layers to have dtype float64 by default, call `tf.keras.backend.set_floatx('float64')`. To change just this layer, pass dtype='float64' to the layer constructor. If you are the author of this layer, you can disable autocasting by passing autocast=False to the base Layer constructor.

[[[[2.86404703e-01 6.57619335e-01]
   [2.97015684e-01 9.82635560e-02]
   [5.12109167e-01 6.61301677e-01]
   [2.73257062e-01 7.16630377e-01]
   [6.89347759e-02 5.56319154e-02]
   [2.68446925e-01 5.54778363e-01]
   [2.86740351e-01 9.18670592e-01]]

  [[1.37923815e-01 8.53850245e-01]
   [2.39049048e-01 4.40782589e-01]
   [5.56441607e-01 1.75341996e-01]
   [2.00587136e-02 4.94617345e-01]
   [7.26282219e-01 3.47326518e-01]
   [1.31355389e-01 2.56608587e-02]
   [7.17480409e-01 5.40043348e-01]]

  [[3.29457013e-01 5.11291441e-01]
   [3.89007326e-01 6.25568050e-01]
   [6.95440981e-01 8.81813478e-01]
   [1.96835611e-02 5.76253816e-01]
   [7.06969026e-04 5.76329261e-01]
   [3.18074238e-01 2.96359529e-01]
   [5.44910042e-01 8.02619339e-01]]

  [[4.14594713e-02 9.04710327e-01]
   [4.85646914e-01 3.95563457e-01]
   [8.34417860e-01 2.93541958e-01]
   [1.37465777e-01 5.91530740e-01]
   [9.01767816e-02 2.10951454e-01]
   [7.57117206e-01 2.83428279e-01]
   [9.92903611e-01 1.13418829e-01]]

  [[4.42290523e-01 5.95948490e-01]
   [7.24276547e-01 8.19404449e-01]
   [7.23147673e-01 4.86682617e-01]
   [1.32073722e-02 4.59669012e-01]
   [4.69288739e-01 2.65081389e-01]
   [5.68099886e-01 5.89725162e-01]
   [8.44584633e-01 9.41648059e-01]]

  [[5.46477097e-01 5.63690576e-01]
   [9.54739437e-01 4.13943212e-01]
   [9.94284810e-01 7.40050707e-01]
   [1.66902099e-03 8.43382213e-01]
   [9.25314686e-01 4.25660493e-01]
   [1.67151242e-01 3.24521094e-02]
   [9.11420177e-01 5.54265451e-01]]

  [[5.73364561e-01 9.09386547e-01]
   [6.82366676e-01 3.41160902e-01]
   [3.05285789e-01 9.76834261e-01]
   [1.36983464e-01 5.10390334e-01]
   [6.89769140e-01 2.23367479e-01]
   [4.47976859e-01 2.85674397e-01]
   [1.41793081e-01 9.01148631e-01]]]]
tensor([[[[0.3284, 0.3346],
          [0.4748, 0.4559]],

         [[0.5549, 0.4440],
          [0.6154, 0.4520]]]])
torch.Size([1, 2, 2, 2])
tensor([[[[0.3284, 0.5549],
          [0.3346, 0.4440]],

         [[0.4748, 0.6154],
          [0.4559, 0.4520]]]])
tf.Tensor(
[[[[0.32842743 0.5549176 ]
   [0.3345996  0.44401568]]

  [[0.4748177  0.6153681 ]
   [0.45593265 0.45198718]]]], shape=(1, 2, 2, 2), dtype=float32)

Process finished with exit code 0

通过对比pytorch的tensor和tensorflow的tensor可以验证改写成功，只不过精度有点区别。

然后就可以定义工具类

from tensorflow.keras import layers
def tensorflow_AdaptiveAvgPool2d(input_size, output_size):
    stridesz = np.floor(input_size / output_size).astype(np.int32)  # 向下取整
    kernelsz = input_size - (output_size - 1) * stridesz
    avg = layers.AveragePooling2D((kernelsz, kernelsz), strides=(stridesz, stridesz))
    return avg

if __name__ == '__main__':
    input_data = np.random.rand(1, 7, 7, 2)
    tensorflow_input = tf.convert_to_tensor(input_data)
    out = tensorflow_AdaptiveAvgPool2d(input_size=7, output_size=2)(tensorflow_input)
    print(out)

ref: https://wenku.baidu.com/view/2dc468fed25abe23482fb4daa58da0116c171f3a.html