用PyTorch实现二维张量上的滑动窗口

import torch

def sliding_window(input_tensor, window_size, stride):
    batch_size, channels, height, width = input_tensor.size()
    output_height = (height - window_size) // stride + 1
    output_width = (width - window_size) // stride + 1

    # 通过unfold函数获取滑动窗口的视图
    unfolded_tensor = input_tensor.unfold(2, window_size, stride).unfold(3, window_size, stride)
    unfolded_tensor = unfolded_tensor.contiguous().view(batch_size, channels, -1, window_size, window_size)

    # 对滑动窗口视图进行操作，这里只是示例，可以根据实际需求进行相应的处理
    output_tensor = unfolded_tensor.mean(dim=(3, 4)) # 取滑动窗口内的均值

    return output_tensor

# 创建一个示例输入张量
input_tensor = torch.arange(1, 41).view(1, 1, 5, 8) # 输入张量的大小为1x1x5x8

# 定义滑动窗口的大小和步长
window_size = 3
stride = 2

# 调用滑动窗口函数
output_tensor = sliding_window(input_tensor, window_size, stride)

print("输入张量：")
print(input_tensor)

print("滑动窗口后的输出张量：")
print(output_tensor)