pytorch学习笔记-激活函数层

90 阅读 0 评论 60 点赞

我是靠谱客的博主疯狂墨镜，这篇文章主要介绍pytorch学习笔记-激活函数层，现在分享给大家，希望可以做个参考。

注意事项：激活函数也是来自nn模块下

复制代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
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()