<reponame>iArunava/scratchai
"""
The GoogLeNet Architecture.
Paper: https://www.cs.unc.edu/~wliu/papers/GoogLeNet.pdf
"""

import torch
import torch.nn as nn
import torch.nn.functional as F

from scratchai.nets.utils import get_net
from scratchai.nets.common import Flatten
from scratchai.pretrained import urls


__all__ = ['InceptionB', 'googlenet', 'googlenet_paper']



def conv(ic:int, oc:int, k, s, p):
  layers = [nn.Conv2d(ic, oc, k, s, p, bias=False), 
            nn.BatchNorm2d(oc, eps=0.001), nn.ReLU(inplace=True)]
  return layers


class AuxClassifier(nn.Sequential):
  """
  The Auxiliary Classifier as suggested by the paper as added on top of
  4a and 4d modules.

  Arguments
  ---------
  ic : int
       The number of in_channels

  nc : int
       The number of classes to output.
  """
  def __init__(self, ic:int, nc:int=1000):
    super().__init__()

    self.net = nn.Sequential(nn.AdaptiveAvgPool2d((4, 4)),
                             *conv(ic, 128, 1, 1, 0), Flatten(),
                             nn.Linear(2048, 1024), nn.Dropout2d(0.7),
                             nn.Linear(1024, nc))
                             

class InceptionB(nn.Module):
  """
  The Inception Module as described in the GoogLeNet paper.

  Arguements
  ----------
  ic                : int
                      The number of in channels
  
  oc1x1             : int
                      The number of out channels from the 1x1 conv

  mc3x3             : int
                      The number of reduce channels from the 1x1 convs 
                      before getting passed into the 3x3 conv.

  oc3x3             : int
                      The number of out channels from the 3x3 conv

  mc5x5             : int
                      The number of reduce channels from the 1x1 convs 
                      before getting passed into the 5x5 conv.

  oc5x5             : int
                      The number of out channels from the 5x5 conv

  ocpool            : int
                      The number of out channels of the 1x1 conv after 
                      the pooling.

  replace5x5with3x3 : bool
                      Changes the branch with 5x5 convs to 3x3 convs
  """
  def __init__(self, ic:int, oc1x1:int, mc3x3, oc3x3, mc5x5, oc5x5, ocpool,
               replace5x5with3x3:bool=False):
    super().__init__()
    
    ks5x5 = 5
    pad5x5 = 2
    if replace5x5with3x3: 
      ks5x5 = 3
      pad5x5 = 1

    self.net1x1 = nn.Sequential(*conv(ic, oc1x1, 1, 1, 0))

    self.net3x3 = nn.Sequential(*conv(ic, mc3x3, 1, 1, 0), 
                                *conv(mc3x3, oc3x3, 3, 1, 1))

    self.net5x5 = nn.Sequential(*conv(ic, mc5x5, 1, 1, 0), 
                                *conv(mc5x5, oc5x5, ks5x5, 1, pad5x5))

    self.pool   = nn.Sequential(nn.MaxPool2d(3, 1, 1, ceil_mode=True), 
                                *conv(ic, ocpool, 1, 1, 0))

  def forward(self, x):
      x1 = self.net1x1(x)
      x2 = self.net3x3(x)
      x3 = self.net5x5(x)
      x4 = self.pool(x)

      x = torch.cat([x1, x2, x3, x4], dim=1)
      return x
    

class GoogLeNet(nn.Module):
  def __init__(self, nc:int=1000, inception=InceptionB, aux:bool=False,
               **kwargs):
    super().__init__()
    self.aux = aux

    self.conv1 = nn.Sequential(*conv(3, 64, 7, 2, 1),
                                nn.MaxPool2d(3, 2, ceil_mode=True))

    self.conv2 = nn.Sequential(*conv(64, 64, 1, 1, 0), 
                               *conv(64, 192, 3, 1, 1),
                               nn.MaxPool2d(3, 2, ceil_mode=True))

    self.inception3a = inception(192, 64, 96, 128, 16, 32, 32, **kwargs)
    self.inception3b = inception(256, 128, 128, 192, 32, 96, 64, **kwargs)

    self.inception4a = inception(480, 192, 96, 208, 16, 48, 64, **kwargs)
    self.inception4b = inception(512, 160, 112, 224, 24, 64, 64, **kwargs)
    self.inception4c = inception(512, 128, 128, 256, 24, 64, 64, **kwargs)
    self.inception4d = inception(512, 112, 144, 288, 32,  64, 64, **kwargs)
    self.inception4e = inception(528, 256, 160, 320, 32, 128, 128, **kwargs)

    self.inception5a = inception(832, 256, 160, 320, 32, 128, 128, **kwargs)
    self.inception5b = inception(832, 384, 192, 384, 48, 128, 128, **kwargs)
    
    if self.aux:
      self.aux1 = AuxClassifier(512, nc)
      self.aux2 = AuxClassifier(528, nc)

    self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
    self.classifier = nn.Sequential(nn.Dropout(0.4), nn.Linear(1024, nc))
    
    self.maxpool = nn.MaxPool2d(3, 2, ceil_mode=True)

  def forward(self, x):
    x = self.conv1(x)
    x = self.conv2(x)

    x = self.inception3a(x)
    x = self.inception3b(x)
    
    x = self.maxpool(x)

    x = self.inception4a(x)
    x = self.inception4b(x)
    x = self.inception4c(x)
    x1 = self.inception4d(x)
    x2 = self.inception4e(x1)
    
    # TODO Add the auxiliary classifiers in the forward pass

    x = self.maxpool(x2)

    x = self.inception5a(x)
    x = self.inception5b(x)
    
    x = self.avgpool(x)
    x = Flatten()(x)
    x = self.classifier(x)
    return x


def googlenet(pretrained=True, **kwargs):
  """
  GoogLeNet Model with weights as given by the officials who 
  trained it on TensorFlow.
  """
  kwargs['aux'] = False if 'aux' not in kwargs else kwargs['aux']
  kwargs['replace5x5with3x3'] = True if 'replace5x5with3x3' not in kwargs \
                                else kwargs['replace5x5with3x3']

  return get_net(GoogLeNet, pretrained=pretrained, fname='googlenet', 
                 kwargs_net=kwargs, attr='classifier', inn=1024,
                 pretrain_url=urls.googlenet_url)


def googlenet_paper(pretrained=False, **kwargs):
  """
  GoogLeNet Model as given in the official Paper.
  """
  kwargs['aux'] = True if 'aux' not in kwargs else kwargs['aux']
  kwargs['replace5x5with3x3'] = False if 'replace5x5with3x3' not in kwargs \
                                else kwargs['replace5x5with3x3']

  return get_net(GoogLeNet, pretrained=pretrained, pretrain_url=None, 
                 fname='googlenet', kwargs_net=kwargs, attr='classifier',
                 inn=1024)

"""
The GoogLeNet Architecture.
Paper: https://www.cs.unc.edu/~wliu/papers/GoogLeNet.pdf
"""

import torch
import torch.nn as nn
import torch.nn.functional as F

from scratchai.nets.utils import get_net
from scratchai.nets.common import Flatten
from scratchai.pretrained import urls


__all__ = ['InceptionB', 'googlenet', 'googlenet_paper']



def conv(ic:int, oc:int, k, s, p):
  layers = [nn.Conv2d(ic, oc, k, s, p, bias=False), 
            nn.BatchNorm2d(oc, eps=0.001), nn.ReLU(inplace=True)]
  return layers


class AuxClassifier(nn.Sequential):
  """
  The Auxiliary Classifier as suggested by the paper as added on top of
  4a and 4d modules.

  Arguments
  ---------
  ic : int
       The number of in_channels

  nc : int
       The number of classes to output.
  """
  def __init__(self, ic:int, nc:int=1000):
    super().__init__()

    self.net = nn.Sequential(nn.AdaptiveAvgPool2d((4, 4)),
                             *conv(ic, 128, 1, 1, 0), Flatten(),
                             nn.Linear(2048, 1024), nn.Dropout2d(0.7),
                             nn.Linear(1024, nc))
                             

class InceptionB(nn.Module):
  """
  The Inception Module as described in the GoogLeNet paper.

  Arguements
  ----------
  ic                : int
                      The number of in channels
  
  oc1x1             : int
                      The number of out channels from the 1x1 conv

  mc3x3             : int
                      The number of reduce channels from the 1x1 convs 
                      before getting passed into the 3x3 conv.

  oc3x3             : int
                      The number of out channels from the 3x3 conv

  mc5x5             : int
                      The number of reduce channels from the 1x1 convs 
                      before getting passed into the 5x5 conv.

  oc5x5             : int
                      The number of out channels from the 5x5 conv

  ocpool            : int
                      The number of out channels of the 1x1 conv after 
                      the pooling.

  replace5x5with3x3 : bool
                      Changes the branch with 5x5 convs to 3x3 convs
  """
  def __init__(self, ic:int, oc1x1:int, mc3x3, oc3x3, mc5x5, oc5x5, ocpool,
               replace5x5with3x3:bool=False):
    super().__init__()
    
    ks5x5 = 5
    pad5x5 = 2
    if replace5x5with3x3: 
      ks5x5 = 3
      pad5x5 = 1

    self.net1x1 = nn.Sequential(*conv(ic, oc1x1, 1, 1, 0))

    self.net3x3 = nn.Sequential(*conv(ic, mc3x3, 1, 1, 0), 
                                *conv(mc3x3, oc3x3, 3, 1, 1))

    self.net5x5 = nn.Sequential(*conv(ic, mc5x5, 1, 1, 0), 
                                *conv(mc5x5, oc5x5, ks5x5, 1, pad5x5))

    self.pool   = nn.Sequential(nn.MaxPool2d(3, 1, 1, ceil_mode=True), 
                                *conv(ic, ocpool, 1, 1, 0))

  def forward(self, x):
      x1 = self.net1x1(x)
      x2 = self.net3x3(x)
      x3 = self.net5x5(x)
      x4 = self.pool(x)

      x = torch.cat([x1, x2, x3, x4], dim=1)
      return x
    

class GoogLeNet(nn.Module):
  def __init__(self, nc:int=1000, inception=InceptionB, aux:bool=False,
               **kwargs):
    super().__init__()
    self.aux = aux

    self.conv1 = nn.Sequential(*conv(3, 64, 7, 2, 1),
                                nn.MaxPool2d(3, 2, ceil_mode=True))

    self.conv2 = nn.Sequential(*conv(64, 64, 1, 1, 0), 
                               *conv(64, 192, 3, 1, 1),
                               nn.MaxPool2d(3, 2, ceil_mode=True))

    self.inception3a = inception(192, 64, 96, 128, 16, 32, 32, **kwargs)
    self.inception3b = inception(256, 128, 128, 192, 32, 96, 64, **kwargs)

    self.inception4a = inception(480, 192, 96, 208, 16, 48, 64, **kwargs)
    self.inception4b = inception(512, 160, 112, 224, 24, 64, 64, **kwargs)
    self.inception4c = inception(512, 128, 128, 256, 24, 64, 64, **kwargs)
    self.inception4d = inception(512, 112, 144, 288, 32,  64, 64, **kwargs)
    self.inception4e = inception(528, 256, 160, 320, 32, 128, 128, **kwargs)

    self.inception5a = inception(832, 256, 160, 320, 32, 128, 128, **kwargs)
    self.inception5b = inception(832, 384, 192, 384, 48, 128, 128, **kwargs)
    
    if self.aux:
      self.aux1 = AuxClassifier(512, nc)
      self.aux2 = AuxClassifier(528, nc)

    self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
    self.classifier = nn.Sequential(nn.Dropout(0.4), nn.Linear(1024, nc))
    
    self.maxpool = nn.MaxPool2d(3, 2, ceil_mode=True)

  def forward(self, x):
    x = self.conv1(x)
    x = self.conv2(x)

    x = self.inception3a(x)
    x = self.inception3b(x)
    
    x = self.maxpool(x)

    x = self.inception4a(x)
    x = self.inception4b(x)
    x = self.inception4c(x)
    x1 = self.inception4d(x)
    x2 = self.inception4e(x1)
    
    # TODO Add the auxiliary classifiers in the forward pass

    x = self.maxpool(x2)

    x = self.inception5a(x)
    x = self.inception5b(x)
    
    x = self.avgpool(x)
    x = Flatten()(x)
    x = self.classifier(x)
    return x


def googlenet(pretrained=True, **kwargs):
  """
  GoogLeNet Model with weights as given by the officials who 
  trained it on TensorFlow.
  """
  kwargs['aux'] = False if 'aux' not in kwargs else kwargs['aux']
  kwargs['replace5x5with3x3'] = True if 'replace5x5with3x3' not in kwargs \
                                else kwargs['replace5x5with3x3']

  return get_net(GoogLeNet, pretrained=pretrained, fname='googlenet', 
                 kwargs_net=kwargs, attr='classifier', inn=1024,
                 pretrain_url=urls.googlenet_url)


def googlenet_paper(pretrained=False, **kwargs):
  """
  GoogLeNet Model as given in the official Paper.
  """
  kwargs['aux'] = True if 'aux' not in kwargs else kwargs['aux']
  kwargs['replace5x5with3x3'] = False if 'replace5x5with3x3' not in kwargs \
                                else kwargs['replace5x5with3x3']

  return get_net(GoogLeNet, pretrained=pretrained, pretrain_url=None, 
                 fname='googlenet', kwargs_net=kwargs, attr='classifier',
                 inn=1024)
The GoogLeNet Architecture. Paper: https://www.cs.unc.edu/~wliu/papers/GoogLeNet.pdf The Auxiliary Classifier as suggested by the paper as added on top of   4a and 4d modules.    Arguments   ---------   ic : int        The number of in_channels    nc : int        The number of classes to output. The Inception Module as described in the GoogLeNet paper.    Arguements   ----------   ic                : int                       The number of in channels      oc1x1             : int                       The number of out channels from the 1x1 conv    mc3x3             : int                       The number of reduce channels from the 1x1 convs                        before getting passed into the 3x3 conv.    oc3x3             : int                       The number of out channels from the 3x3 conv    mc5x5             : int                       The number of reduce channels from the 1x1 convs                        before getting passed into the 5x5 conv.    oc5x5             : int                       The number of out channels from the 5x5 conv    ocpool            : int                       The number of out channels of the 1x1 conv after                        the pooling.    replace5x5with3x3 : bool                       Changes the branch with 5x5 convs to 3x3 convs # TODO Add the auxiliary classifiers in the forward pass GoogLeNet Model with weights as given by the officials who    trained it on TensorFlow. GoogLeNet Model as given in the official Paper.