| import torch |
| import torch.nn as nn |
| from huggingface_hub import PyTorchModelHubMixin |
|
|
|
|
| class ODCNN(nn.Module, PyTorchModelHubMixin): |
| def __init__(self, dropout_rate=0.5): |
| super().__init__() |
|
|
| |
| |
| self.conv1 = nn.Conv2d(3, 10, kernel_size=(3, 7)) |
| self.relu1 = nn.ReLU() |
| self.pool1 = nn.MaxPool2d(kernel_size=(3, 1), stride=(3, 1)) |
|
|
| |
| self.conv2 = nn.Conv2d(10, 20, kernel_size=(3, 3)) |
| self.relu2 = nn.ReLU() |
| self.pool2 = nn.MaxPool2d(kernel_size=(3, 1), stride=(3, 1)) |
|
|
| |
| |
| self.flatten_size = 20 * 8 * 7 |
|
|
| |
| self.dropout = nn.Dropout(p=dropout_rate) |
|
|
| |
| self.fc1 = nn.Linear(self.flatten_size, 256) |
| self.relu_fc = nn.ReLU() |
|
|
| |
| self.fc2 = nn.Linear(256, 1) |
| self.sigmoid = nn.Sigmoid() |
|
|
| def forward(self, x): |
| x = self.conv1(x) |
| x = self.relu1(x) |
| x = self.pool1(x) |
|
|
| x = self.conv2(x) |
| x = self.relu2(x) |
| x = self.pool2(x) |
|
|
| x = x.view(x.size(0), -1) |
|
|
| x = self.dropout(x) |
| x = self.fc1(x) |
| x = self.relu_fc(x) |
|
|
| x = self.dropout(x) |
| x = self.fc2(x) |
| x = self.sigmoid(x) |
|
|
| return x |
|
|
|
|
| if __name__ == "__main__": |
| from torchinfo import summary |
|
|
| model = ODCNN() |
| summary(model, (1, 3, 80, 15)) |
|
|
