import torch
import torchvision
from tensorboardX import SummaryWriter
from torch import nn
from torch.nn import Conv2d
from torch.utils.data import DataLoader
dataset = torchvision.datasets.CIFAR10('./dataset',train=False,transform=torchvision.transforms.ToTensor(),download=True)
dataloader = DataLoader(dataset,batch_size=64)
class Tudui(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = Conv2d(3,6,kernel_size=3, stride=1, padding=0)
#3是Inputchannel,6是outputchannel
def forward(self,x):
x = self.conv1(x)
return x
tudui = Tudui()
writer = SummaryWriter('logs_conv')
step = 0
for data in dataloader:
imgs,targets = data
output = tudui(imgs)
# print(imgs.shape)
# print(output.shape)
#返回torch.Size([64,3,32,32]) batchsize64,3通道,32*32
# 返回torch.Size([64,6,30,30]) Batchsize64,卷积后变6通道,30*30
writer.add_images('input',imgs,step)
output = torch.reshape(output,(-1,3,30,30))
#writer里面添加Image需要是三通道的,卷积之后是6通道,所以在这里reshape成三通道,-1是自适应有多大的batchsize
writer.add_images('output',output,step)
step += 1
writer.close()
【PyTorch笔记】pytorch入门教程7卷积层
于 2023-05-13 20:31:12 首次发布