Update pages/15_CNN.py

pages/15_CNN.py

@@ -1,4 +1,4 @@
-import streamlit as st
+]import streamlit as st
 import torch
 import torch.nn as nn
 import torch.optim as optim
@@ -47,20 +47,31 @@ class CNN(nn.Module):
             nn.BatchNorm2d(64),
             nn.ReLU(),
             nn.MaxPool2d(2))
-        self.fc1 = nn.Linear(6*6*64, 600)
+        
+        # Automatically determine the size of the flattened features after convolution and pooling
+        self._to_linear = None
+        self.convs(torch.randn(1, 3, 32, 32))
+        
+        self.fc1 = nn.Linear(self._to_linear, 600)
         self.drop = nn.Dropout2d(0.25)
         self.fc2 = nn.Linear(600, 100)
         self.fc3 = nn.Linear(100, 10)
 
+    def convs(self, x):
+        x = self.layer1(x)
+        x = self.layer2(x)
+        if self._to_linear is None:
+            self._to_linear = x.view(x.size(0), -1).shape[1]
+        return x
+    
     def forward(self, x):
-        out = self.layer1(x)
-        out = self.layer2(out)
-        out = out.view(out.size(0), -1)
-        out = self.fc1(out)
-        out = self.drop(out)
-        out = self.fc2(out)
-        out = self.fc3(out)
-        return out
+        x = self.convs(x)
+        x = x.view(x.size(0), -1)
+        x = self.fc1(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.fc3(x)
+        return x
 
 model = CNN().to(device)