Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| class FFTCNN(nn.Module): | |
| """ | |
| Defines the Convolutional Neural Network architecture. | |
| This structure must match the model that was trained and saved. | |
| """ | |
| def __init__(self): | |
| super(FFTCNN, self).__init__() | |
| # Ensure 'self.' is used here to define the layers as instance attributes | |
| self.conv_layers = nn.Sequential( | |
| nn.Conv2d(1, 16, kernel_size=3, padding=1), | |
| nn.ReLU(), | |
| nn.MaxPool2d(kernel_size=2, stride=2), | |
| nn.Conv2d(16, 32, kernel_size=3, padding=1), | |
| nn.ReLU(), | |
| nn.MaxPool2d(kernel_size=2, stride=2) | |
| ) | |
| # Ensure 'self.' is used here as well | |
| self.fc_layers = nn.Sequential( | |
| nn.Linear(32 * 56 * 56, 128), # This size depends on your 224x224 input | |
| nn.ReLU(), | |
| nn.Linear(128, 2) # 2 output classes | |
| ) | |
| def forward(self, x): | |
| # Now, 'self.conv_layers' can be found because it was defined correctly | |
| x = self.conv_layers(x) | |
| x = x.view(x.size(0), -1) # Flatten the feature maps | |
| x = self.fc_layers(x) | |
| return x |