text: ## TODO: define the convolutional neural network architecture
  
  import torch
  from torch.autograd import Variable
  import torch.nn as nn
  import torch.nn.functional as F
  # can use the below import should you choose to initialize the weights of your Net
  import torch.nn.init as I
  
  
  class NaimishNet(nn.Module):
  
      def __init__(self):
          super(Net, self).__init__()
          # Based on Facial Key Points Detection using Deep Convolutional Neural Network - NaimishNet
          # https://arxiv.org/pdf/1710.00977.pdf
          self.conv1 = nn.Conv2d(1, 32, 4)
          self.conv2 = nn.Conv2d(32, 64, 3)
          self.conv3 = nn.Conv2d(64, 128, 2)
          self.conv4 = nn.Conv2d(128, 256, 1)
          
          self.pool = nn.MaxPool2d(2, 2)
          
          # Normalization layer after every pooling layer.
          self.bn1 = nn.BatchNorm2d(32)
          self.bn2 = nn.BatchNorm2d(64)
          self.bn3 = nn.BatchNorm2d(128)
          self.bn4 = nn.BatchNorm2d(256)
          
          self.dropout1 = nn.Dropout(p=0.1)
          self.dropout2 = nn.Dropout(p=0.2)
          self.dropout3 = nn.Dropout(p=0.3)
          self.dropout4 = nn.Dropout(p=0.4)
          self.dropout5 = nn.Dropout(p=0.5)
          self.dropout6 = nn.Dropout(p=0.6)
      
          self.fc1 = nn.Linear(43264, 1000)
          self.fc2 = nn.Linear(1000, 136)        
  
          
      def forward(self, x):
          x = self.pool(F.relu(self.conv1(x)))
          x = self.bn1(x)
          x = self.dropout1(x)
  
          x = self.pool(F.relu(self.conv2(x)))
          x = self.bn2(x)
          x = self.dropout2(x)
  
          x = self.pool(F.relu(self.conv3(x)))
          x = self.bn3(x)
          x = self.dropout3(x)
  
          x = self.pool(F.relu(self.conv4(x)))
          x = self.bn4(x)
          x = self.dropout4(x)
        
          # Flatten
          x = x.view(x.size(0), -1)
          
          x = self.dropout5(x)
          x = F.relu(self.fc1(x))
          x = self.dropout6(x) 
          x = self.fc2(x)
          
        
          return x
  
  
  from torchvision.models.resnet import conv3x3, BasicBlock, Bottleneck
  
  # Minor modification of ResNet18 in torchvision
  # - Removed layers 3 & 4
  class ResFace(nn.Module):
      def __init__(self, zero_init_residual=False):
          super(ResFace, self).__init__()
          
          self.inplanes = 224
          block = BasicBlock
          self.conv1 = nn.Conv2d(1, 224, kernel_size=7, stride=2, padding=3,
                                 bias=False)
          self.bn1 = nn.BatchNorm2d(224)
          self.relu = nn.ReLU(inplace=True)
          self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
          self.layer1 = self._make_layer(block, 64, 2)
          self.layer2 = self._make_layer(block, 128, 2, stride=2)
          self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
          self.fc = nn.Linear(128 * block.expansion, 136)
  
          for m in self.modules():
              if isinstance(m, nn.Conv2d):
                  nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
              elif isinstance(m, nn.BatchNorm2d):
                  nn.init.constant_(m.weight, 1)
                  nn.init.constant_(m.bias, 0)
  
          # Zero-initialize the last BN in each residual branch,
          # so that the residual branch starts with zeros, and each residual block behaves like an identity.
          # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
          if zero_init_residual:
              for m in self.modules():
                  if isinstance(m, Bottleneck):
                      nn.init.constant_(m.bn3.weight, 0)
                  elif isinstance(m, BasicBlock):
                      nn.init.constant_(m.bn2.weight, 0)
  
      def _make_layer(self, block, planes, blocks, stride=1):
          downsample = None
          if stride != 1 or self.inplanes != planes * block.expansion:
              downsample = nn.Sequential(
                  nn.Conv2d(self.inplanes, planes * block.expansion, kernel_size=1, stride=stride, bias=False),
                  nn.BatchNorm2d(planes * block.expansion),
              )
  
          layers = []
          layers.append(block(self.inplanes, planes, stride, downsample))
          self.inplanes = planes * block.expansion
          for _ in range(1, blocks):
              layers.append(block(self.inplanes, planes))
  
          return nn.Sequential(*layers)
  
      def forward(self, x):
          x = self.conv1(x)
          x = self.bn1(x)
          x = self.relu(x)
          x = self.maxpool(x)
  
          x = self.layer1(x)
          x = self.layer2(x)
  
          x = self.avgpool(x)
          x = x.view(x.size(0), -1)
          x = self.fc(x)
  
          return x
  
  