【Pytorch实现】——最大池化和平均池化

import torch
from torch import nn

def pool2d(X, pool_size, mode='max'):
  # 获取池化层核的大小
  p_h, p_w = pool_size
  # 计算经过最大池化后的特征图的大小
  Y = torch.zeros((X.shape[0] - p_h + 1,X.shape[1] - p_w + 1))
  # 行列遍历进行最大池化
  for i in range(Y.shape[0]):
    for j in range(Y.shape[1]):
      if mode == 'max':
        Y[i,j] = X[i:i + p_h, j:j + p_w].max()
      elif mode == 'avg':
        Y[i,j] = X[i:i + p_h, j:j + p_w].mean()

  return Y

X = torch.tensor([[0.0,1.0,2.0],[3.0,4.0,5.0],[6.0,7.0,8.0]])
pool2d(X,(2,2))