import torch
import torchvision.datasets
from torch import nn
from torch.nn import Conv2d
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

dataset = torchvision.datasets.CIFAR10("./DXAA",train=False,transform=torchvision.transforms.ToTensor(),download=True)#路径为返回上一层路径，建立DATA文件夹
# 注意这里的dataset 需要将其转成tensor的类型，否则下面的卷积操作将无法进行
dataloader = DataLoader(dataset,batch_size=64)

class sza_(nn.Module):#继承神经网络的模型Module
    def __init__(self):
        super(sza_, self).__init__()
        self.conv1 = Conv2d(in_channels=3,out_channels=6,kernel_size=3,stride=1,padding=0)#彩色图片为3，黑白图片为1,kernel_size 为数字，那么卷积核为3*3
        #现在我们的网络当中有一个卷积层叫做conv1
    def forward(self,x):
        x = self.conv1(x)#这里的输入值为为tensor类型。
        return x

sza__ = sza_()#实例化、

# sza_(
#   (conv1): Conv2d(3, 6, kernel_size=(3, 3), stride=(1, 1))
# )
writer = SummaryWriter("./logs")
step = 0
for data in dataloader:
    imgs,targets = data
    output = sza__(imgs)
    # print(imgs.shape)
    # print(output.shape)
# torch.Size([64, 3, 32, 32])意思为imgs 的batchsize为64一组  ，in_channels =3, 大小是32*32的
# torch.Size([64, 6, 30, 30])经过一轮卷积之后 output的batchsize仍然为64一组，但是out_channels 输出为6， 大小缩水为30*30
    writer.add_images("卷积操作before",imgs,step)
    output = torch.reshape(output,(-1,3,30,30))#torch.size(64,6,30,30)--->(xxx,3,30,30)
    # print(output.shape)#reshape结果：torch.Size([128, 3, 30, 30])
    writer.add_images("卷积操作after",output,step)#    writer.add_images("卷积操作after",output,step)这里显示报错原因在于写入的图片out_channels为6，而彩色图片的channels为3
    step = step + 1
writer.close()

效果如下：