Update app.py

app.py

@@ -3,24 +3,35 @@ from fastai.vision.all import *
 from PIL import Image
 #
 #learn = load_learner('export.pkl')
-learn = torch.load('digit_classifier.pth')
-learn.eval() #switch to eval mode
+#learn = torch.load('digit_classifier.pth')
+#learn.eval() #switch to eval mode
+model_dict=torch.load('my_model.pt')
+W1,B1,W2,B2,W3,B3=model_dict['W1'],model_dict['B1'],model_dict['W2'],model_dict['B2'],model_dict['W3'],model_dict['B3']
+def mdlV2(xb):
+    res = xb@W1+B1
+    res = res.max(tensor(0.))
+    res = res@W2+B2  # returns 10 features for each input
+    res = res.max(tensor(0.))
+    res = res@W3+B3  # returns 10 features for each input
+    return res
+
+
 labels = [str(x) for x in range(10)]
-#################################
-#Define class for importing Model
-class DigitClassifier(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.fc1 = torch.nn.Linear(64, 32)
-        self.fc2 = torch.nn.Linear(32, 16)
-        self.fc3 = torch.nn.Linear(16, 10)
+# #################################
+# #Define class for importing Model
+# class DigitClassifier(torch.nn.Module):
+#     def __init__(self):
+#         super().__init__()
+#         self.fc1 = torch.nn.Linear(64, 32)
+#         self.fc2 = torch.nn.Linear(32, 16)
+#         self.fc3 = torch.nn.Linear(16, 10)
 
-    def forward(self, x):
-        x = x.view(-1, 64)
-        x = torch.relu(self.fc1(x))
-        x = torch.relu(self.fc2(x))
-        x = self.fc3(x)
-        return x
+#     def forward(self, x):
+#         x = x.view(-1, 64)
+#         x = torch.relu(self.fc1(x))
+#         x = torch.relu(self.fc2(x))
+#         x = self.fc3(x)
+#         return x
 #########################################
 #Define function to reduce image of arbitrary size to 8x8 per model requirements.
 def reduce_image_count(image):
@@ -45,7 +56,7 @@ def predict(img):
     pic = np.array(gray_img) #convert to array
     inp_img=reduce_image_count(pic)#Reduce image to required input size
     
-    otpt=F.softmax(learn.forward(inp_img.view(-1,64)))
+    otpt=F.softmax(mdlV2(inp_img.view(-1,64)))
     #pred,pred_idx,probs = learn.predict(img)
     
     return dict([[labels[i], float(otpt[0].data[i])] for i in range(len(labels))]),inp_img