I built this cnn model for image classification ( cats and dogs ) and and my model presents overffiting … I don’t know how to fix this problem
class CNNModel(nn.Module):
def __init__(self):
super(CNNModel, self).__init__()
ch = [3, 64, 128, 224, 256, 394, 512, 1024, 2048, 4096]
kwargs = {'kernel_size': 5, 'stride': 1, 'padding': 2, 'bias': False}
self.conv1 = nn.Sequential( # (3, 64, 64) -->> (3, 224, 224)
nn.Conv2d(ch[0], ch[3], **kwargs), # (3, 64, 64) -->> (3, 64, 64)
nn.ReLU(), # (64, 64, 64) -->> (224, 64, 64)
nn.BatchNorm2d(ch[3]),
nn.MaxPool2d(kernel_size=2) # (64, 32, 32) -->> (224, 32, 32)
)
self.conv2 = nn.Sequential( # (64, 32, 32) -->> (224, 32, 32)
nn.Conv2d(ch[3], ch[5], **kwargs), # (128, 32, 32) -->> (394, 32, 32)
nn.ReLU(), # (128, 32, 32) -->> (394, 32, 32)
nn.BatchNorm2d(ch[5]),
)
self.conv3 = nn.Sequential( # (128, 32, 32) -->> (394, 32, 32)
nn.Conv2d(ch[5], ch[6], **kwargs), # (256, 32, 32) -->> (512, 32, 32)
nn.ReLU(), # (256, 32, 32) -->> (512, 32, 32)
nn.BatchNorm2d(ch[6]),
nn.MaxPool2d(kernel_size=2) # (256, 16, 16) -->> (512, 16, 16)
)
self.conv4 = nn.Sequential( # (256, 16, 16) -->> (512, 16, 16)
nn.Conv2d(ch[6], ch[7], **kwargs), # (394, 16, 16) -->> (1024, 16, 16)
nn.BatchNorm2d(ch[7]),
nn.ReLU(), # (394, 16, 16) -->> (1024, 16, 16)
)
self.conv5 = nn.Sequential( # (394, 16, 16) -->> (1024, 16, 16)
nn.Conv2d(ch[7], ch[8], **kwargs), # (512, 16, 16) -->> (2048, 16, 16)
nn.BatchNorm2d(ch[8]),
nn.ReLU(), # (512, 16, 16) -->> (2048, 16, 16)
nn.MaxPool2d(kernel_size=2) # (512, 8, 8) -->> (2048, 8, 8)
)
self.fc1 = nn.Sequential(
nn.Linear(2048 * 8 * 8, 1024),
nn.ReLU(),
nn.Dropout(0.5)
)
self.fc2 = nn.Sequential(
nn.Linear(1024, 1024),
nn.ReLU(),
nn.Dropout(0.6)
)
self.fc3 = nn.Sequential(
nn.Linear(1024, 12),
nn.ReLU(),
nn.Dropout(0.5)
)
self.softmax = nn.LogSoftmax()
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
x = self.conv3(x)
x = self.conv4(x)
x = self.conv5(x)
x = x.view(x.size(0), -1)
x = self.fc1(x)
x = self.fc2(x)
x = self.fc3(x)
output = self.softmax(x)
return output