MindSpore实现AdaptiveAvgPool2d

这是昇腾AI创新大赛2022-昇思赛道参赛踩坑记录的第二篇，上一篇主要讲了 MindSpore 和 PyTorch 常用算子的 API 映射。

本文主要纪录如何在 MindSpore 中实现 AdaptiveAvgPool2d。以下是官方 API 文档：
mindspore.nn.AdaptiveAvgPool2d

（在写这篇文章的时候发现 MindSpore 已经更新到 1.8 版本了，真快啊！）

在 MindSpore1.6 版本之前没有提供 AdaptiveAvgPool2d，从 1.6 版本开始提供，但不支持 Ascend，只支持 GPU。没办法，参赛必须使用 Ascend，所以自己写。

1、什么是 AdaptiveAvgPool2d？

要复现首先就要知道什么是 AdaptiveAvgPool2d，AdaptiveAvgPool2d包含以下几个概念：
二元（2d）
池化（Pool）
均值（Avg）
自适应（Adaptive）

前3个比较简单，可以参考下这个回答：Pytorch 里 nn.AdaptiveAvgPool2d(output_size) 原理是什么?

我们主要来讲讲自适应（Adaptive）部分。

这里存在一个误区：AdaptiveAvgPool2d 的功能就是在 AvgPool2d 的基础上加上了自动计算池化操作时的kernel_size 和 stride 等数据，用户只需要传入 “待处理数据” 和 “目标大小” 就可以获得目标输出。基于上述理解，我们只需要通过 input_size 和 output_size 反推出 kernel_size、stride 等参数，就可以使用 AvgPool2d 代替了。

但如果真的这样做会发现计算结果误差很大，原因在于出发点就错了，AdaptiveAvgPool2d 底层的实现方法与上述方法不一致。

那么什么是 AdaptiveAvgPool2d，可以参考：What is AdaptiveAvgPool2d?

这里主要记录一下 PyTorch 的实现方法以及如何用 MindSpore 去复现。

2、PyTorch AdaptiveAvgPool2d

import torch

def torch_AdaptiveAvgPool2d(inputs, target_size):
    """ NCHW """
    H = target_size[0]
    W = target_size[1]

    H_start = (torch.arange(H, dtype=torch.float32) * (inputs.size(-2) / H)).long()
    H_end = ((torch.arange(H, dtype=torch.float32)+1) * (inputs.size(-2) / H)).ceil().long()

    W_start = (torch.arange(W, dtype=torch.float32) * (inputs.size(-1) / W)).long()
    W_end = ((torch.arange(W, dtype=torch.float32)+1) * (inputs.size(-1) / W)).ceil().long()

    pooled2 = []
    for idx_H in range(H):
        pooled1 = []
        for idx_W in range(W):
            pooled1.append(torch.mean(inputs[:, :, H_start[idx_H]:H_end[idx_H], W_start[idx_W]:W_end[idx_W]], dim=(-2,-1), keepdim=True))
        pooled1 = torch.cat(pooled1, -1)
        pooled2.append(pooled1)
    pooled2 = torch.cat(pooled2,-2)

    return pooled2

data = [[[[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]]
,[[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]]
]]

# 输入
inputs = torch.tensor(data, dtype=torch.float32)
print("input_size:", inputs.size())
print(inputs)
print("---------------------------------------------------")

# 测试torch_AdaptiveAvgPool2d
avgpool1 = torch_AdaptiveAvgPool2d(inputs, (2,3))
print("avgpool1_shape:", avgpool1.shape)
print(avgpool1)
print("---------------------------------------------------")

# 与torch.nn.AdaptiveAvgPool2d对比
avgpool2 = torch.nn.AdaptiveAvgPool2d((2,3))(inputs)
print("avgpool2_shape:", avgpool2.shape)
print(avgpool2)

# Out
"""
input_size: torch.Size([1, 2, 6, 8])
tensor([[[[2., 3., 4., 5., 6., 9., 7., 8.],
          [2., 3., 4., 5., 6., 9., 7., 8.],
          [2., 3., 4., 5., 6., 9., 7., 8.],
          [2., 3., 4., 5., 6., 9., 7., 8.],
          [2., 3., 4., 5., 6., 9., 7., 8.],
          [2., 3., 4., 5., 6., 9., 7., 8.]],

         [[2., 3., 4., 5., 6., 9., 7., 8.],
          [2., 3., 4., 5., 6., 9., 7., 8.],
          [2., 3., 4., 5., 6., 9., 7., 8.],
          [2., 3., 4., 5., 6., 9., 7., 8.],
          [2., 3., 4., 5., 6., 9., 7., 8.],
          [2., 3., 4., 5., 6., 9., 7., 8.]]]])
---------------------------------------------------
avgpool1_shape: torch.Size([1, 2, 2, 3])
tensor([[[[3., 6., 8.],
          [3., 6., 8.]],

         [[3., 6., 8.],
          [3., 6., 8.]]]])
---------------------------------------------------
avgpool2_shape: torch.Size([1, 2, 2, 3])
tensor([[[[3., 6., 8.],
          [3., 6., 8.]],

         [[3., 6., 8.],
          [3., 6., 8.]]]])
"""

3、MindSpore AdaptiveAvgPool2d

而 MindSpore 的代码，只要替换掉上述 PyTorch 代码的 torch.arange， torch.mean， torch.cat，ceil 和 long 就可以了。

torch.arange --> mindspore.numpy.arange
torch.mean --> mindspore.ops.ReduceMean
torch.cat --> mindspore.ops.Concat
ceil --> mindspore.numpy.ceil 或者 mindspore.ops.Ceil (ops.Ceil 只支持 Ascend)
long --> mindspore.ops.Cast(x, dtype=mstype.int64)

import mindspore.numpy as np
from mindspore import dtype as mstype
from mindspore import ops, Tensor, context, nn

class AdaptiveAvgPool2d(nn.Cell):
    def __init__(self, output_size):
        """Initialize AdaptiveAvgPool2d."""
        super(AdaptiveAvgPool2d, self).__init__()
        self.output_size = output_size

    def adaptive_avgpool2d(self, inputs):
        """ NCHW """
        H = self.output_size[0]
        W = self.output_size[1]

        H_start = ops.Cast()(np.arange(start=0, stop=H, dtype=mstype.float32) * (inputs.shape[-2] / H), mstype.int64)
        H_end = ops.Cast()(np.ceil(((np.arange(start=0, stop=H, dtype=mstype.float32)+1) * (inputs.shape[-2] / H))), mstype.int64)

        W_start = ops.Cast()(np.arange(start=0, stop=W, dtype=mstype.float32) * (inputs.shape[-1] / W), mstype.int64)
        W_end = ops.Cast()(np.ceil(((np.arange(start=0, stop=W, dtype=mstype.float32)+1) * (inputs.shape[-1] / W))), mstype.int64)

        pooled2 = []
        for idx_H in range(H):
            pooled1 = []
            for idx_W in range(W):
                h_s = int(H_start[idx_H].asnumpy())
                h_e = int(H_end[idx_H].asnumpy())
                w_s = int(W_start[idx_W].asnumpy())
                w_e = int(W_end[idx_W].asnumpy())
                res = inputs[:, :, h_s:h_e, w_s:w_e]
                # res = inputs[:, :, H_start[idx_H]:H_end[idx_H], W_start[idx_W]:W_end[idx_W]]  # 这样写mindspore tensor切片报类型错误，不知道为啥
                pooled1.append(ops.ReduceMean(keep_dims=True)(res, (-2,-1)))
            pooled1 = ops.Concat(-1)(pooled1)
            pooled2.append(pooled1)
        pooled2 = ops.Concat(-2)(pooled2)

        return pooled2

    def construct(self, x):
        x = self.adaptive_avgpool2d(x)
        return x


data = [[[[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]]
,[[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]]
]]

# 输入
inputs = Tensor(data, dtype=mstype.float32)
print("input_size:", inputs.shape)
print(inputs)
print("---------------------------------------------------")

# 测试mindspore_AdaptiveAvgPool2d
avgpool = AdaptiveAvgPool2d(output_size=(2,3))
output = avgpool(inputs)
print("ms_avgpool_shape:", output.shape)
print(output)

# Out
"""
input_size: (1, 2, 6, 8)
[[[[2. 3. 4. 5. 6. 9. 7. 8.]
   [2. 3. 4. 5. 6. 9. 7. 8.]
   [2. 3. 4. 5. 6. 9. 7. 8.]
   [2. 3. 4. 5. 6. 9. 7. 8.]
   [2. 3. 4. 5. 6. 9. 7. 8.]
   [2. 3. 4. 5. 6. 9. 7. 8.]]

  [[2. 3. 4. 5. 6. 9. 7. 8.]
   [2. 3. 4. 5. 6. 9. 7. 8.]
   [2. 3. 4. 5. 6. 9. 7. 8.]
   [2. 3. 4. 5. 6. 9. 7. 8.]
   [2. 3. 4. 5. 6. 9. 7. 8.]
   [2. 3. 4. 5. 6. 9. 7. 8.]]]]
---------------------------------------------------
ms_avgpool_shape: (1, 2, 2, 3)
[[[[3. 6. 8.]
   [3. 6. 8.]]

  [[3. 6. 8.]
   [3. 6. 8.]]]]
"""

可以看到，输出和 torch.nn.AdaptiveAvgPool2d 一致，迁移成功。

在复现的过程中发现 MindSpore 使用

 res = inputs[:, :, H_start[idx_H]:H_end[idx_H], W_start[idx_W]:W_end[idx_W]] 

对 Tensor 做切片操作时会发生数据类型方面的异常，其中 H_start[idx_H]，H_end[idx_H]，W_start[idx_W] 和 W_end[idx_W]] 为 Tensor.int64 类型，不清楚是什么原因导致类型匹配出错。

所以我将上述四个 Tensor 变量转为 numpy 再转为 int，然后再进行切片，解决了报错。

---

参考：

mindspore 版 AdaptiveAvgPool2d 替代方案参考