Implementation of CNN/ConvNet Model using PyTorch (depicted in the picture above)
class CNN(torch.nn.Module):
def __init__(self):
    super(CNN, self).__init__()
    # L1 ImgIn shape=(?, 28, 28, 1)
    #    Conv     -> (?, 28, 28, 32)
    #    Pool     -> (?, 14, 14, 32)
    self.layer1 = torch.nn.Sequential(
        torch.nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1),
        torch.nn.ReLU(),
        torch.nn.MaxPool2d(kernel_size=2, stride=2),
        torch.nn.Dropout(p=1 - keep_prob))
    # L2 ImgIn shape=(?, 14, 14, 32)
    #    Conv      ->(?, 14, 14, 64)
    #    Pool      ->(?, 7, 7, 64)
    self.layer2 = torch.nn.Sequential(
        torch.nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),
        torch.nn.ReLU(),
        torch.nn.MaxPool2d(kernel_size=2, stride=2),
        torch.nn.Dropout(p=1 - keep_prob))
    # L3 ImgIn shape=(?, 7, 7, 64)
    #    Conv      ->(?, 7, 7, 128)
    #    Pool      ->(?, 4, 4, 128)
    self.layer3 = torch.nn.Sequential(
        torch.nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
        torch.nn.ReLU(),
        torch.nn.MaxPool2d(kernel_size=2, stride=2, padding=1),
        torch.nn.Dropout(p=1 - keep_prob))

    # L4 FC 4x4x128 inputs -> 625 outputs
    self.fc1 = torch.nn.Linear(4 * 4 * 8192, 240, bias=True)
    torch.nn.init.xavier_uniform(self.fc1.weight)
    self.layer4 = torch.nn.Sequential(
        self.fc1,
        torch.nn.ReLU(),
        torch.nn.Dropout(p=1 - keep_prob))
    # L5 Final FC 625 inputs -> 10 outputs
    self.fc2 = torch.nn.Linear(240, 2, bias=True)
    torch.nn.init.xavier_uniform_(self.fc2.weight) # initialize parameters

def forward(self, x):
    out = self.layer1(x)
    out = self.layer2(out)
    out = self.layer3(out)
    out = out.view(out.size(0), -1)   # Flatten them for FC
    out = self.fc1(out)
    out = self.fc2(out)
    return out

instantiate CNN model
model = CNN()
model