Spaces:
Sleeping
Sleeping
File size: 1,201 Bytes
b7c5baf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 |
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 |