Upload 5 files

Model.py

@@ -0,0 +1,20 @@
+import torch
+
+class MNIST(torch.nn.Module):
+    def __init__(self):
+        super(MNIST, self).__init__()
+        self.conv = torch.nn.Sequential(torch.nn.Conv2d(1, 32, 3, 1, 1),
+                                        torch.nn.ReLU(),
+                                        torch.nn.Conv2d(32, 64, 3, 1, 1),
+                                        torch.nn.ReLU(),
+                                        torch.nn.MaxPool2d(2, 2))
+        self.dense = torch.nn.Sequential(torch.nn.Linear(14 * 14 * 64, 1024),
+                                         torch.nn.ReLU(),
+                                         torch.nn.Dropout(p=0.2),
+                                         torch.nn.Linear(1024, 10))
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = x.view(-1, 14 * 14 * 64)
+        x = self.dense(x)
+        return x

inference.py

@@ -0,0 +1,25 @@
+import cv2
+import torch
+import torchvision.transforms as transforms
+import time
+from Model import MNIST
+
+
+def images2tensor(image):
+    img = cv2.imread(image)
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    transf = transforms.ToTensor()
+    img_tensor = torch.unsqueeze(transf(img), dim=0)
+    return img_tensor
+
+
+if __name__ == "__main__":
+    start = time.time()
+    device = torch.device('cpu')
+    model = MNIST().to(device)
+    model.load_state_dict(torch.load('mnist.pkl'))  # load
+    input_data = images2tensor("0.png")
+    res = model(input_data)
+    end = time.time()
+    print("手写数字图片检测的结果为：", res.argmax())
+    print("infer time: ", end - start)

mnist.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f23ce66930c4c38a18e1b51ab4306e2c70319bb3ce7ef1820bf3390876e0b751
+size 51503055

requirements.txt

@@ -0,0 +1,4 @@
+torch
+torchvision
+opencv-python
+opencv-python-headless