| import torch |
| import torch.nn as nn |
| import torch.optim as optim |
|
|
| class CNNModel(nn.Module): |
|
|
| def __init__(self): |
| super(CNNModel, self).__init__() |
| self.conv1 = nn.Conv2d(3, 64, kernel_size=5, padding=2) |
| self.relu1 = nn.ReLU() |
| self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2) |
| self.conv2 = nn.Conv2d(64, 128, kernel_size=5, padding=2) |
| self.relu2 = nn.ReLU() |
| self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2) |
| self.conv3 = nn.Conv2d(128, 256, kernel_size=3, padding=1) |
| self.relu3 = nn.ReLU() |
| self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2) |
| self.conv4 = nn.Conv2d(256, 384, kernel_size=5, padding=1) |
| self.relu4 = nn.ReLU() |
| self.conv5 = nn.Conv2d(384, 256, kernel_size=1, padding=0) |
| self.relu5 = nn.ReLU() |
| self.avgpool = nn.AdaptiveAvgPool2d((3, 3)) |
| self.flatten = nn.Flatten() |
| self.dropout1 = nn.Dropout(p=0.1) |
| self.fc1 = nn.Linear(256 * 3 * 3, 1024) |
| self.relu6 = nn.ReLU() |
| self.dropout2 = nn.Dropout(p=0.1) |
| self.fc2 = nn.Linear(1024, 512) |
| self.relu7 = nn.ReLU() |
| self.fc3 = nn.Linear(512, 200) |
| self.softmax = nn.Softmax(dim=1) |
|
|
| def forward(self, x): |
| x = self.pool1(self.relu1(self.conv1(x))) |
| x = self.pool2(self.relu2(self.conv2(x))) |
| x = self.pool3(self.relu3(self.conv3(x))) |
| x = self.relu4(self.conv4(x)) |
| x = self.relu5(self.conv5(x)) |
| x = self.avgpool(x) |
| x = self.flatten(x) |
| x = self.dropout1(x) |
| x = self.relu6(self.fc1(x)) |
| x = self.dropout2(x) |
| x = self.relu7(self.fc2(x)) |
| x = self.fc3(x) |
| x = self.softmax(x) |
| return x |
| |
| model = CNNModel() |
| loss_fn = nn.NLLLoss() |
| optimizer = optim.Adam(model.parameters(), lr=0.0003) |
