phenolicat commited on
Commit
bf5cc8d
1 Parent(s): 456d283

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +15 -12
app.py CHANGED
@@ -13,25 +13,28 @@ st.write("")
13
  # enable users to upload images for the model to make predictions
14
  file_up = st.file_uploader("Upload an image")
15
 
16
- # Define a more complex linear layer
17
- class ComplexLinearLayer(nn.Module):
18
- def __init__(self, input_size, output_size):
19
- super(ComplexLinearLayer, self).__init__()
20
- self.fc1 = nn.Linear(input_size, 512)
21
- self.fc2 = nn.Linear(512, 256)
22
- self.fc3 = nn.Linear(256, output_size)
 
23
 
24
  def forward(self, x):
 
 
 
 
25
  x = x.view(x.size(0), -1)
26
  x = F.relu(self.fc1(x))
27
- x = F.relu(self.fc2(x))
28
- x = self.fc3(x)
29
  x = F.log_softmax(x, dim=1)
30
  return x
31
 
32
- # Create an instance of the model with the more complex linear layer
33
- model = ComplexLinearLayer(input_size=28*28, output_size=10)
34
-
35
  def predict(image):
36
 
37
  # load the model
 
13
  # enable users to upload images for the model to make predictions
14
  file_up = st.file_uploader("Upload an image")
15
 
16
+ # Define a more complex neural network model
17
+ class ComplexModel(nn.Module):
18
+ def __init__(self, input_channels, output_size):
19
+ super(ComplexModel, self).__init__()
20
+ self.conv1 = nn.Conv2d(input_channels, 32, kernel_size=3, stride=1, padding=1)
21
+ self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
22
+ self.fc1 = nn.Linear(64 * 7 * 7, 128)
23
+ self.fc2 = nn.Linear(128, output_size)
24
 
25
  def forward(self, x):
26
+ x = F.relu(self.conv1(x))
27
+ x = F.max_pool2d(x, kernel_size=2, stride=2)
28
+ x = F.relu(self.conv2(x))
29
+ x = F.max_pool2d(x, kernel_size=2, stride=2)
30
  x = x.view(x.size(0), -1)
31
  x = F.relu(self.fc1(x))
32
+ x = self.fc2(x)
 
33
  x = F.log_softmax(x, dim=1)
34
  return x
35
 
36
+ model = ComplexModel(1, 10)
37
+
 
38
  def predict(image):
39
 
40
  # load the model