File size: 2,135 Bytes
c248b9b |
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 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
import torch
import torch.nn as nn
import torch.nn.functional as F
class ExcitometerModel(nn.Module):
def __init__(self, num_classes=10):
super(ExcitometerModel, self).__init__()
# Extended architecture
self.conv1 = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=3, stride=1, padding=1)
self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=1)
self.conv3 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1)
self.conv4 = nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1)
self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
self.dropout = nn.Dropout(p=0.5) # Dropout with 50% probability
self.fc1 = nn.Linear(128 * 8 * 8, 1024) # Adjust dimensions based on input size
self.fc2 = nn.Linear(1024, 512)
self.fc3 = nn.Linear(512, num_classes)
self.batch_norm1 = nn.BatchNorm2d(16)
self.batch_norm2 = nn.BatchNorm2d(32)
self.batch_norm3 = nn.BatchNorm2d(64)
self.batch_norm4 = nn.BatchNorm2d(128)
def forward(self, x):
# Forward pass through the network
x = F.relu(self.batch_norm1(self.conv1(x)))
x = self.pool(F.relu(self.batch_norm2(self.conv2(x))))
x = F.relu(self.batch_norm3(self.conv3(x)))
x = self.pool(F.relu(self.batch_norm4(self.conv4(x))))
x = F.adaptive_avg_pool2d(x, (1, 1)) # Global Average Pooling
x = x.view(-1, 128) # Flatten for the fully connected layer
x = F.relu(self.fc1(x))
x = self.dropout(x) # Apply dropout
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
# Example usage
if __name__ == "__main__":
model = ExcitometerModel(num_classes=10)
print(model)
# Example input tensor with batch size 1, 1 channel, and 64x64 spatial dimensions
example_input = torch.randn(1, 1, 64, 64)
output = model(example_input)
print(output)
|