pytorch用MaxPooling实现图片转换到指定大小

参考自：

• 使用Icecream提高Python调试效率
• Pytorch Tensor交换维度
• MAXPOOL2D

import os
from icecream import ic
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image

path = "./1.jpeg"

# 图片本身
img = Image.open(path)
img.save("old.jpg")
# img.show()

img = torch.from_numpy(np.array(img, dtype=np.float32))
ic(img.shape)  # torch.Size([680, 440,3])

# 添加模型需要的batch_size，将channel换到第二个维度
img = img.unsqueeze(0).transpose(1, 3).transpose(2, 3)
ic(img.shape)  # torch.Size([1, 3, 680, 440])


# 先随便写池化层的kernel_size和步长
m = nn.MaxPool2d((3, 2), stride=(2, 1))

new_img = m(img)
ic(new_img.shape)  # torch.Size([1, 3,  339, 439])

# 维度转变 成Image需要的
new_img = new_img.transpose(1, 2).transpose(2, 3).squeeze()
ic(new_img.shape)  # torch.Size([339, 439, 3])

img = Image.fromarray(np.uint8(new_img)).convert('RGB')
img.save("./new.jpg")

其中新旧图片的尺寸和步长、kernel_size的计算公式如下:
$\text{ Input: }\left(N,C,H_{in},W_{in}\right)\text{ or }\left(C,H_{in},W_{in}\right)$
$\text{ Output: }\left(N,C,H_{\text{out }},W_{\text{out }}\right)\text{ or }\left(C,H_{\text{out }},W_{\text{out }}\right)\text{, where }$

效果如图