GbrlOl commited on
Commit
704b95d
1 Parent(s): 4431339

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +52 -0
app.py CHANGED
@@ -1,4 +1,56 @@
1
  import gradio as gr
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2
 
3
  def predict(im):
4
 
 
1
  import gradio as gr
2
+ import torch
3
+ from torch import nn
4
+ from torch.utils.data import DataLoader
5
+ from torchvision import datasets
6
+ from torchvision.transforms import ToTensor
7
+ import torch.nn.functional as F
8
+
9
+ device = (
10
+ "cuda"
11
+ if torch.cuda.is_available()
12
+ else "mps"
13
+ if torch.backends.mps.is_available()
14
+ else "cpu"
15
+ )
16
+
17
+ class CNN(nn.Module):
18
+ def __init__(self):
19
+ super(CNN, self).__init__()
20
+ # Definimos las capas convolucionales
21
+ self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, padding=1)
22
+ self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1)
23
+ self.conv3 = nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1)
24
+
25
+ # Definimos capas fully connected
26
+ self.fc1 = nn.Linear(128 * 3 * 3, 256)
27
+ self.fc2 = nn.Linear(256, 10)
28
+
29
+ # Definimos un max pooling y dropout
30
+ self.pool = nn.MaxPool2d(2, 2)
31
+ self.dropout = nn.Dropout(0.25)
32
+
33
+ def forward(self, x):
34
+ # Pasamos las entradas por las capas convolucionales y el max pooling
35
+ x = self.pool(F.relu(self.conv1(x)))
36
+ x = self.pool(F.relu(self.conv2(x)))
37
+ x = self.pool(F.relu(self.conv3(x)))
38
+
39
+ # Aplanamos la salida de las capas convolucionales para pasar a fully connected
40
+ x = x.view(-1, 128 * 3 * 3)
41
+
42
+ # Pasamos por las capas fully connected
43
+ x = F.relu(self.fc1(x))
44
+ x = self.dropout(x)
45
+ x = self.fc2(x)
46
+
47
+ return x
48
+
49
+ model = CNN().to(device)
50
+
51
+ # Cargar el modelo en la CPU
52
+ model = CNN().to(device)
53
+ model.load_state_dict(torch.load("model_mnist_cnn.pth", map_location=torch.device('cpu')))
54
 
55
  def predict(im):
56