注意事项:激活函数也是来自nn模块下
import torchvision
from tensorboardX import SummaryWriter
from torch import nn
from torch.nn import ReLU
from torch.utils.data import DataLoader
from torchvision import transforms
# input = torch.tensor([[1,-0.5],
# [-1,3]])
#
# output = torch.reshape(input,(-1,1,2,2))
# print(output.shape)
train_set = torchvision.datasets.CIFAR10(root="datasets2",train=False,transform=transforms.ToTensor(),download=True)
train_loader = DataLoader(dataset=train_set,batch_size=64)
class zj_relu(nn.Module):
def __init__(self) -> None:
super().__init__()
self.relu = ReLU()
def forward(self,input):
output = self.relu(input)
return output
zj_relu_1 = zj_relu()
# output = zj_relu_1(input)
# print(output)
step = 0
writer = SummaryWriter("logs")
for data in train_loader:
imgs,targets = data
output = zj_relu_1(imgs)
writer.add_images("input",imgs, step)
writer.add_images("output",output,step)
step += 1
writer.close()