Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| class SeizureDetector(nn.Module): | |
| def __init__(self, num_classes=2): | |
| super(SeizureDetector, self).__init__() | |
| self.conv1= nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1) # 32, 224, 224 | |
| self.pool= nn.MaxPool2d(kernel_size=2, stride=2) # 32, 112, 112 | |
| self.conv2= nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1) # 64, 112, 112 -> 64, 56, 56 | |
| self.conv3= nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1) # 128, 56, 56 -> 128, 28, 28 | |
| self.conv4= nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1) # 256, 28, 28 -> 256, 14, 14 | |
| # Adding Batch Normalization | |
| self.bn1 = nn.BatchNorm2d(32) | |
| self.bn2 = nn.BatchNorm2d(64) | |
| self.bn3 = nn.BatchNorm2d(128) | |
| self.bn4 = nn.BatchNorm2d(256) | |
| self.dropout = nn.Dropout(p=0.5) # Dropout with a probability of 50% | |
| self.fc1= nn.Linear(256*14*14, 120) | |
| self.fc2= nn.Linear(120, 32) | |
| self.fc3= nn.Linear(32, num_classes) | |
| def forward(self, x): | |
| x = self.pool(F.relu(self.bn1(self.conv1(x)))) # 32, 112, 112 | |
| x = self.pool(F.relu(self.bn2(self.conv2(x)))) # 64, 56, 56 | |
| x = self.pool(F.relu(self.bn3(self.conv3(x)))) # 128, 28, 28 | |
| x = self.pool(F.relu(self.bn4(self.conv4(x)))) # 256, 14, 14 | |
| x = torch.flatten(x, 1) | |
| x = self.dropout(F.relu(self.fc1(x))) # Apply dropout | |
| x = self.dropout(F.relu(self.fc2(x))) # Apply dropout | |
| x = self.fc3(x) | |
| return x |