轻松学pytorch-构建卷积神经网络
轻松学pytorch-构建卷积神经网络
网络实现
大家好，这个是我的pytorch学习笔记第三篇，主要是使用pytorch来构建一个简单的卷积神经网络，完成mnist手写字符识别。涉及到主要知识点是如何使用torch.nn.Module这个基类来实现一个网络结构定义。这个基类中最重要的是实现自己的forward方法，这个也是自定义网络结构的实现方法。我实现的简单CNN_Net类的代码如下：
class CNN_Net(t.nn.Module):
def __init__(self):
super(CNN_Net, self).__init__()
self.cnn_layers = t.nn.Sequential(
t.nn.Conv2d(in_channels=1, out_channels=8, kernel_size=3, padding=1, stride=1),
t.nn.MaxPool2d(kernel_size=2, stride=2),
t.nn.ReLU(),
t.nn.Conv2d(in_channels=8, out_channels=32, kernel_size=3, padding=1, stride=1),
t.nn.MaxPool2d(kernel_size=2, stride=2),
t.nn.ReLU(),
)
self.fc_layers = t.nn.Sequential(
t.nn.Linear(7 * 7 * 32, 200),
t.nn.ReLU(),
t.nn.Linear(200, 100),
t.nn.ReLU(),
t.nn.Linear(100, 10),
t.nn.LogSoftmax(dim=1)
)
def forward(self, x):
out = self.cnn_layers(x)
out = out.view(-1, 7 * 7 * 32)
out = self.fc_layers(out)
return out这其中有几个需要说明的地方
Torch.nn.Conv2d中默认的填充方式valid，就是padding=0，如果你想处理边缘像素，需要手动指定填充像素宽度。
- in_channels表示输入通道数目
- out_channels表示输出通道数目
在从卷积层到全链接层过渡中，flatten需要自己实现一下，如果不想实现，一个偷懒的方法就是直接写死，我就是采用view方法来实现flatten的。
训练与测试
基于交叉熵损失完成了训练，对模型进行eval之后就可以调用跟保存模型了，另外说一下为什么要对训练好的模型进行eval，eval的作用是对模型训练状态下的一些层在测试或者推理阶段是不需要的，所以可以eval来适当调整与简化模型。训练与测试代码如下：
model = CNN_Net().cuda()
loss_fn= t.nn.CrossEntropyLoss()
optimizer = t.optim.Adam(model.parameters(), lr=1e-3)
for s in range(5):
print("run in step : %d"%s)
for i, (x_train, y_train) in enumerate(train_dl):
x_train = x_train.cuda()
y_train = y_train.cuda()
y_pred = model.forward(x_train)
train_loss = loss_fn(y_pred, y_train)
if (i + 1) % 100 == 0:
print(i + 1, train_loss.item())
optimizer.zero_grad()
train_loss.backward()
optimizer.step()
total = 0;
correct_count = 0
model.eval()
for test_images, test_labels in test_dl:
with t.no_grad():
pred_labels = model(test_images.cuda())
predicted = t.max(pred_labels, 1)[1]
correct_count += (predicted == test_labels.cuda()).sum()
total += len(test_labels)
print(correct_count.cpu().detach().numpy(), total)
print("total acc : %.4f\n"%(correct_count.cpu().detach().numpy() / total))
t.save(model, './cnn_mnist_model.pt')OpenCV调用模型
保存了模型之后，还可以转化为ONNX格式，把模型送给OpenCV DNN模块调用，这块我也做了个简单的测试，发现预测良好！
首先需要把模型转化为ONNX格式，这个可以使用torch中工具包，几行代码即可搞定，请看我的转化代码
dummy_input = torch.randn(1, 1, 28, 28, device='cuda')
model = torch.load("./cnn_mnist_model.pt")
torch.onnx.export(model, dummy_input, "cnn_mnist.onnx", verbose=True)这样我就得到了onnx的文件，直接通过OpenCV DNN模块加载调用试试
import cv2 as cv
import numpy as np
mnist_net = cv.dnn.readNetFromONNX("cnn_mnist.onnx")
image = cv.imread("D:/images/9_99.png")
gray = cv.cvtColor(image, cv.COLOR_BGR2GRAY)
cv.imshow("input", gray)
blob = cv.dnn.blobFromImage(gray, 0.00392, (28, 28), (0.5)) / 0.5
print(blob.shape)
mnist_net.setInput(blob)
result = mnist_net.forward()
pred_label = np.argmax(result, 1)
print("predit label : %d"%pred_label)
cv.waitKey(0)
cv.destroyAllWindows()这个其中值得注意的是，你之前预处理数据是做totensor跟normalization，所以opencv读取的图片也要做一下，我加载是mnist测试集中的数字9图像，预测结果如下：
完全没问题。靠谱！
通过这个例子，我完成了模型实现，训练与保存，ONNX转化，第三方OpenCV DNN使用。还是有很多值得实践的地方，希望跟大家一起继续进步！请大家点赞支持！我继续坚持写下去！