"""
    构建一个具有 池化层 的神经网络
        -- 相当于图像的下采样，降低图像的分辨率
"""
import torchvision
from torch import nn

from torch.nn import Conv2d, MaxPool2d
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter("logs/011")

dataset = torchvision.datasets.CIFAR10(root="./visionData", train=False, transform=torchvision.transforms.ToTensor(),
                                       download=True)
dataloader = DataLoader(dataset, batch_size=64)


class MyNn(nn.Module):
    def __init__(self):
        super(MyNn, self).__init__()
        '''
            kernel_size: 卷积核尺寸
            ceil_mode: 卷积核移动到边缘且未完全覆盖时，是否保留结果
        '''
        self.maxpool = MaxPool2d(kernel_size=3, ceil_mode=True)

    def forward(self, input):
        output = self.maxpool(input)
        return output


nnInstance = MyNn()
step = 0
for data in dataloader:
    imgs, targets = data
    writer.add_images("011input", imgs, step)

    output = nnInstance(imgs)
    writer.add_images("011output", output, step)

    step = step + 1

writer.close()