class CNNMnist(nn.Module):
def init(self, args):
super(CNNMnist, self).init()
self.conv1 = nn.Conv2d(args.num_channels, 10, kernel_size=5)
self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
self.conv2_drop = nn.Dropout2d()
self.fc1 = nn.Linear(320, 50)
self.fc2 = nn.Linear(50, args.num_classes)
def forward(self, x):
x = F.relu(F.max_pool2d(self.conv1(x), 2))
x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
x = x.view(-1, x.shape[1]*x.shape[2]*x.shape[3])
x = F.relu(self.fc1(x))
x = F.dropout(x, training=self.training)
x = self.fc2(x)
return F.log_softmax(x, dim=1)
conv2_drop is a Dropout layer. Dropout is a regularization technique used to prevent overfitting in neural networks. During training, the Dropout layer randomly sets a portion of its input units to 0, which helps make the model more robust and not overly reliant on any single input in the training set. In this case, nn.Dropout2d() is a two-dimensional Dropout, which randomly sets entire channels to 0, which is particularly useful in convolutional neural networks. This means that during training, the entire feature map (after the conv2 convolutional layer) will be randomly turned off. This helps prevent overfitting and enhances the model's generalization ability. During testing, the Dropout layer does not have any effect, and all neurons are used.
x = x.view(-1, x.shape[1]*x.shape[2]*x.shape[3]): Then, the shape of x is changed (or flattened) so that it can be processed by the fully connected layer. This step is commonly referred to as flattening.
x = F.dropout(x, training=self.training): Then, x goes through a Dropout layer, which helps prevent overfitting. Note that the Dropout layer only works during training, and during testing, all neurons are used.
The input is an image with num_channels, and the output is num_classes similarity scores (converted to probability density), with the class with the highest similarity being selected as the result.