| import torch |
| import torch.nn as nn |
|
|
| class CatDogClassifier(nn.Module): |
| def __init__(self): |
| super(CatDogClassifier, self).__init__() |
| |
| |
| self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1) |
| self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1) |
| self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1) |
| |
| |
| self.pool = nn.MaxPool2d(2, 2) |
| self.relu = nn.ReLU() |
| self.dropout = nn.Dropout(0.5) |
| |
| |
| self.fc1 = nn.Linear(128 * 28 * 28, 512) |
| self.fc2 = nn.Linear(512, 128) |
| self.fc3 = nn.Linear(128, 2) |
| |
| def forward(self, x): |
| x = self.relu(self.conv1(x)) |
| x = self.pool(x) |
| x = self.relu(self.conv2(x)) |
| x = self.pool(x) |
| x = self.relu(self.conv3(x)) |
| x = self.pool(x) |
| x = x.view(x.size(0), -1) |
| x = self.dropout(self.relu(self.fc1(x))) |
| x = self.dropout(self.relu(self.fc2(x))) |
| x = self.fc3(x) |
| return x |
|
|
