Update README.md

README.md

@@ -22,97 +22,31 @@ This gradio app predicts digits using a convolutive neural network (CNN) that wa
 }
 
 The PyTorch network architecture:
-CNN(
-  (conv1): Sequential(
-    (0): Conv2d(1, 16, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
-    (1): ReLU()
-    (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
-  )
-  (conv2): Sequential(
-    (0): Conv2d(16, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
-    (1): ReLU()
-    (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
-  )
-  (out): Linear(in_features=1568, out_features=10, bias=True)
+Sequential(
+  (conv1): Conv2d(1, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
+  (relu1): ReLU()
+  (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
+  (conv2): Conv2d(32, 64, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
+  (relu2): ReLU()
+  (pool2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
+  (flatten): Flatten(start_dim=1, end_dim=-1)
+  (fc1): Linear(in_features=3136, out_features=1024, bias=True)
+  (dropout): Dropout(p=0.5, inplace=False)
+  (fc2): Linear(in_features=1024, out_features=10, bias=True)
 )
 
-The model was trained using cross entropy loss function, the Adam stochastic optimizer.  Training was done for 5 epochs, using batch size of 64, and with a learning rate of 0.001.  The loss and training accuracy record was as follows:
-
-Epoch 1
--------------------------------
-loss: 2.308622  [    0/60000]
-loss: 0.263496  [ 6400/60000]
-loss: 0.153133  [12800/60000]
-loss: 0.174262  [19200/60000]
-loss: 0.065622  [25600/60000]
-loss: 0.139496  [32000/60000]
-loss: 0.064824  [38400/60000]
-loss: 0.040185  [44800/60000]
-loss: 0.147770  [51200/60000]
-loss: 0.132756  [57600/60000]
-Test Error: 
- Accuracy: 97.8%, Avg loss: 0.068824 
-
-Epoch 2
--------------------------------
-loss: 0.081690  [    0/60000]
-loss: 0.140447  [ 6400/60000]
-loss: 0.024672  [12800/60000]
-loss: 0.021052  [19200/60000]
-loss: 0.093074  [25600/60000]
-loss: 0.057293  [32000/60000]
-loss: 0.059690  [38400/60000]
-loss: 0.041043  [44800/60000]
-loss: 0.008444  [51200/60000]
-loss: 0.023947  [57600/60000]
-Test Error: 
- Accuracy: 98.8%, Avg loss: 0.037573 
-
-Epoch 3
--------------------------------
-loss: 0.026763  [    0/60000]
-loss: 0.008322  [ 6400/60000]
-loss: 0.004648  [12800/60000]
-loss: 0.022180  [19200/60000]
-loss: 0.026884  [25600/60000]
-loss: 0.127719  [32000/60000]
-loss: 0.029921  [38400/60000]
-loss: 0.018033  [44800/60000]
-loss: 0.020504  [51200/60000]
-loss: 0.015732  [57600/60000]
-Test Error: 
- Accuracy: 98.5%, Avg loss: 0.044061 
-
-Epoch 4
--------------------------------
-loss: 0.001018  [    0/60000]
-loss: 0.047208  [ 6400/60000]
-loss: 0.015693  [12800/60000]
-loss: 0.025309  [19200/60000]
-loss: 0.008418  [25600/60000]
-loss: 0.047959  [32000/60000]
-loss: 0.003479  [38400/60000]
-loss: 0.049097  [44800/60000]
-loss: 0.009243  [51200/60000]
-loss: 0.033639  [57600/60000]
-Test Error: 
- Accuracy: 98.8%, Avg loss: 0.036810 
-
-Epoch 5
--------------------------------
-loss: 0.002880  [    0/60000]
-loss: 0.011252  [ 6400/60000]
-loss: 0.102632  [12800/60000]
-loss: 0.045496  [19200/60000]
-loss: 0.007643  [25600/60000]
-loss: 0.003103  [32000/60000]
-loss: 0.093026  [38400/60000]
-loss: 0.022969  [44800/60000]
-loss: 0.027332  [51200/60000]
-loss: 0.005514  [57600/60000]
-Test Error: 
- Accuracy: 99.0%, Avg loss: 0.027122 
-
-Done!
-
-
+The model was trained using cross entropy loss function, the Adam stochastic optimizer.
+Training was done for 10 epochs, using batch size of 64, and with a learning rate of 0.001.
+
+The training and validation accuracy after each epoch was as follows:
+
+ Epoch 1 ---- train accuracy: 0.9607 ---- val accuracy: 0.9885
+ Epoch 2 ---- train accuracy: 0.9870 ---- val accuracy: 0.9894
+ Epoch 3 ---- train accuracy: 0.9893 ---- val accuracy: 0.9892
+ Epoch 4 ---- train accuracy: 0.9925 ---- val accuracy: 0.9914
+ Epoch 5 ---- train accuracy: 0.9932 ---- val accuracy: 0.9872
+ Epoch 6 ---- train accuracy: 0.9947 ---- val accuracy: 0.9904
+ Epoch 7 ---- train accuracy: 0.9946 ---- val accuracy: 0.9888
+ Epoch 8 ---- train accuracy: 0.9947 ---- val accuracy: 0.9920
+ Epoch 9 ---- train accuracy: 0.9962 ---- val accuracy: 0.9912
+ Epoch 10 ---- train accuracy: 0.9960 ---- val accuracy: 0.9915