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 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65
| import torch from torch import nn from matplotlib import pyplot as plt
class LeNet5(nn.Module): ''' LeNet5网络 INPUT -> 图像规格(28, 28, 1), 待分类数(10) ''' def __init__(self): super(LeNet5, self).__init__() self.conv1 = nn.Sequential( nn.Conv2d(1, 6, kernel_size=5, padding=2, bias=False), nn.Tanh(), nn.MaxPool2d(kernel_size=2, stride=2) ) self.conv2 = nn.Sequential( nn.Conv2d(6, 16, kernel_size=5, bias=False), nn.Tanh(), nn.MaxPool2d(kernel_size=2, stride=2) ) self.classifier = nn.Sequential( nn.Linear(16*5*5, 120), nn.Tanh(), nn.Linear(120, 84), nn.Tanh(), nn.Linear(84, 10) )
def forward(self, x): x = self.conv1(x) x = self.conv2(x) x = x.view(x.size(0), 16*5*5) x = self.classifier(x) return x
net = LeNet5() net.load_state_dict(torch.load("./model/net_best.pth"))
weights = net.conv2._modules['0'].weight.data w = weights.numpy()
fig=plt.figure(figsize=(20, 8)) columns = 5 rows = 2 for i in range(0, columns*rows): fig.add_subplot(rows, columns, i+1) plt.imshow(w[i][0], cmap='gray') print('Second convolutional layer') plt.show()
|