Spaces:
Runtime error
Runtime error
added new model, moved net to separate file, updated app to use transformed data
c7d1130
import torch | |
import torch.nn as nn | |
import torch.nn.functional as F | |
class Net(nn.Module): | |
def __init__(self): | |
super().__init__() | |
self.conv1 = nn.Conv2d(3, 32, 5) | |
self.conv2 = nn.Conv2d(32, 64, 5) | |
self.conv3 = nn.Conv2d(64, 128, 5) | |
x = torch.randn(3, 150, 150).view(-1, 3, 150, 150) | |
self._to_linear = None | |
self.convs(x) | |
self.fc1 = nn.Linear(self._to_linear, 512) | |
self.fc2 = nn.Linear(512, 6) | |
def convs(self, x): | |
x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2)) | |
x = F.max_pool2d(F.relu(self.conv2(x)), (2, 2)) | |
x = F.max_pool2d(F.relu(self.conv3(x)), (2, 2)) | |
if self._to_linear is None: | |
self._to_linear = x[0].shape[0]*x[0].shape[1]*x[0].shape[2] | |
return x | |
def forward(self, x): | |
x = self.convs(x) | |
x = x.view(-1, self._to_linear) | |
x = F.relu(self.fc1(x)) | |
x = self.fc2(x) | |
return F.softmax(x, dim=1) | |