训练、验证代码逻辑
L
e
N
e
t
−
M
o
d
e
l
(
p
y
t
o
r
c
h
版
本
)
LeNet-Model(pytorch版本)
LeNet−Model(pytorch版本)
import torch.nn as nn
import torch.nn.functional as F
import torch
class LeNet(nn.Module):
def __init__(self, num_classes,num_linear=44944):
super(LeNet, self).__init__()
self.conv1 = nn.Conv2d(3, 6, 5)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(num_linear, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, num_classes)
def forward(self, x):
out = F.relu(self.conv1(x))
out = F.max_pool2d(out, 2)
out = F.relu(self.conv2(out))
out = F.max_pool2d(out, 2)
out = out.view(out.size(0), -1)
out = F.relu(self.fc1(out))
out = F.relu(self.fc2(out))
out = self.fc3(out)
return out
def initialize_weights(self):
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.xavier_normal_(m.weight.data)
if m.bias is not None:
m.bias.data.zero_()
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
elif isinstance(m, nn.Linear):
nn.init.normal_(m.weight.data, 0, 0.1)
m.bias.data.zero_()
rgb = torch.randn(1, 3, 224, 224)
net = LeNet(num_classes=8,num_linear=44944)
out = net(rgb)
print('-----'*5)
print(out.shape)
print('-----'*5)
![在这里插入图片描述](https://i-blog.csdnimg.cn/blog_migrate/c2aa244ee884480b7c9271cb79d7f846.png#pic_center)