symbol_inception-resnet-v1.R

# Inception resnet v1, suitable for images with around 299 x 299
#
# Reference:
# Szegedy C, Ioffe S, Vanhoucke V. Inception-v4, inception-resnet and 
# the impact of residual connections on learning, 2016.
# Link to the paper: https://arxiv.org/abs/1602.07261
#
library(mxnet)

Conv <- function(data, num_filter, kernel=c(1, 1), stride=c(1, 1), pad=c(0, 0), 
                 name, suffix="", withRelu=TRUE, withBn=FALSE){
  conv <- mx.symbol.Convolution(data=data, num_filter=num_filter, kernel=kernel, 
                               stride=stride, pad=pad,
                               name=paste0(name, suffix, "_conv2d"))
  if (withBn){
    conv <- mx.symbol.BatchNorm(data=conv, name=paste0(name, suffix, "_bn"))
  }
  if (withRelu){
    conv <- mx.symbol.Activation(data=conv, act_type="relu", 
                                name=paste0(name, suffix, "_relu"))
  }
  
  return(conv)
}

# Input Shape is 299*299*3 (th)
InceptionResnetStem <- function(data,
                                num_1_1, num_1_2, num_1_3,
                                num_2_1, num_2_2, num_2_3,
                                name){
  stem_3x3 <- Conv(data=data, num_filter=num_1_1, kernel=c(3, 3), stride=c(2, 2),
                  name=paste0(name, "_conv"))
  stem_3x3 <- Conv(data=stem_3x3, num_filter=num_1_2, kernel=c(3, 3), 
                  name=paste0(name, "_stem"), suffix="_conv_1")
  stem_3x3 <- Conv(data=stem_3x3, num_filter=num_1_3, kernel=c(3, 3), 
                  name=paste0(name, "_stem"), suffix="_conv_2")
  pool1 <- mx.symbol.Pooling(data=stem_3x3, kernel=c(3, 3), stride=c(2, 2), 
                            pool_type="max", name=paste0("max_", name, "_pool1"))
  stem_1_3x3 <- Conv(data=pool1, num_filter=num_2_1, name=paste0(name, "_stem_1"), 
                    suffix="_conv_1")
  stem_1_3x3 <- Conv(data=stem_1_3x3, num_filter=num_2_2, kernel=c(3, 3), 
                    name=paste0(name, "_stem_1"), suffix="_conv_2")
  stem_1_3x3 <- Conv(data=stem_1_3x3, num_filter=num_2_3, kernel=c(3, 3), 
                    pad=c(1, 1), stride=c(2, 2), name=paste0(name, "_stem_1"), 
                    suffix="_conv_3", withRelu=FALSE)
  bn1 <- mx.symbol.BatchNorm(data=stem_1_3x3, name=paste0(name, "_bn1"))
  act1 <- mx.symbol.Activation(data=bn1, act_type="relu", name=paste0(name, "_relu1"))
  
  return(act1)
}

InceptionResnetA <- function(data,
                             num_1_1,
                             num_2_1, num_2_2,
                             num_3_1, num_3_2, num_3_3,
                             proj,
                             name,
                             scaleResidual=TRUE){
  init <- data
  
  a1 <- Conv(data=data, num_filter=num_1_1, name=paste0(name, "_a_1"), suffix="_conv")
  
  a2 <- Conv(data=data, num_filter=num_2_1, name=paste0(name, "_a_2"), suffix="_conv_1")
  a2 <- Conv(data=a2, num_filter=num_2_2, kernel=c(3, 3), pad=c(1, 1), 
            name=paste0(name, "_a_2"), suffix="_conv_2")
  
  a3 <- Conv(data=data, num_filter=num_3_1, name=paste0(name, "_a_3"), suffix="_conv_1")
  a3 <- Conv(data=a3, num_filter=num_3_2, kernel=c(3, 3), pad=c(1, 1), 
            name=paste0(name, "_a_3"), suffix="_conv_2")
  a3 <- Conv(data=a3, num_filter=num_3_3, kernel=c(3, 3), pad=c(1, 1), 
            name=paste0(name, "_a_3"), suffix="_conv_3")
  # concat
  merge_lst <- list()
  merge_lst <- c(a1, a2, a3)
  merge_lst$num.args <- length(merge_lst)
  merge_lst$name <- paste0(name, "_a_concat1")
  merge <- mxnet:::mx.varg.symbol.Concat(merge_lst)
  
  conv <- Conv(data=merge, num_filter=proj, name=paste0(name, "_a_liner_conv"), 
              withRelu=FALSE)
  if(scaleResidual){
    conv <- conv*0.1
  }
  
  out <- init + conv
  bn <- mx.symbol.BatchNorm(data=out, name=paste0(name, "_a_bn1"))
  act <- mx.symbol.Activation(data=bn, act_type="relu", name=paste0(name, "_a_relu1"))
  
  return(act)
}

InceptionResnetB <- function(data,
                             num_1_1,
                             num_2_1, num_2_2, num_2_3,
                             proj,
                             name,
                             scaleResidual=TRUE){
  init <- data
  
  b1 <- Conv(data=data, num_filter=num_1_1, name=paste0(name, "_b_1"), suffix="_conv")
  
  b2 <- Conv(data=data, num_filter=num_2_1, name=paste0(name, "_b_2"), suffix="_conv_1")
  b2 <- Conv(data=b2, num_filter=num_2_2, kernel=c(1, 7), pad=c(0, 3), 
            name=paste0(name, "_b_2"), suffix="_conv_2")
  b2 <- Conv(data=b2, num_filter=num_2_3, kernel=c(7, 1), pad=c(3, 0), 
            name=paste0(name, "_b_2"), suffix="_conv_3")
  
  merge_lst <- list()
  merge_lst <- c(b1, b2)
  merge_lst$num.args <- length(merge_lst)
  merge_lst$name <- paste0(name, "_b_concat1")
  merge <- mxnet:::mx.varg.symbol.Concat(merge_lst)
  
  conv <- Conv(data=merge, num_filter=proj, name=paste0(name, "_b_liner_conv"), 
              withRelu=FALSE)
  if(scaleResidual){
    conv <- conv*0.1
  }
  
  out <- init + conv
  bn <- mx.symbol.BatchNorm(data=out, name=paste0(name, "_b_bn1"))
  act <- mx.symbol.Activation(data=bn, act_type="relu", name=paste0(name, "_b_relu1"))
  
  return(act)
}

InceptionResnetC <- function(data,
                             num_1_1,
                             num_2_1, num_2_2, num_2_3,
                             proj,
                             name,
                             scaleResidual=TRUE){
  
  init <- data
  
  c1 <- Conv(data=data, num_filter=num_1_1, name=paste0(name, "_c_1"), suffix="_conv")
  
  c2 <- Conv(data=data, num_filter=num_2_1, name=paste0(name, "_c_2"), suffix="_conv_1")
  c2 <- Conv(data=c2, num_filter=num_2_2, kernel=c(1, 3), pad=c(0, 1), 
            name=paste0(name, "_c_2"), suffix="_conv_2")
  c2 <- Conv(data=c2, num_filter=num_2_3, kernel=c(3, 1), pad=c(1, 0), 
            name=paste0(name, "_c_2"), suffix="_conv_3")
  
  merge_lst <- list()
  merge_lst <- c(c1, c2)
  merge_lst$num.args <- length(merge_lst)
  merge_lst$name <- paste0(name, "_c_concat1")
  merge <- mxnet:::mx.varg.symbol.Concat(merge_lst)
  
  conv <- Conv(data=merge, num_filter=proj, name=paste0(name, "_b_liner_conv"), 
              withRelu=FALSE)
  if(scaleResidual){
    conv <- conv*0.1
  }
  
  out <- init + conv
  bn <- mx.symbol.BatchNorm(data=out, name=paste0(name, "_c_bn1"))
  act <- mx.symbol.Activation(data=bn, act_type="relu", name=paste0(name, "_c_relu1"))
  
  return(act)
}

ReductionResnetA <- function(data,
                             num_2_1,
                             num_3_1, num_3_2, num_3_3,
                             name){
  
  ra1 <- mx.symbol.Pooling(data=data, kernel=c(3, 3), stride=c(2, 2), 
                          pool_type="max", name=paste0("max_", name, "_pool1"))
  
  ra2 <- Conv(data=data, num_filter=num_2_1, kernel=c(3, 3), stride=c(2, 2), 
             name=paste0(name, "_ra_2"), suffix="_conv", withRelu=FALSE)
  
  ra3 <- Conv(data=data, num_filter=num_3_1, name=paste0(name, "_ra_3"), suffix="_conv_1")
  ra3 <- Conv(data=ra3, num_filter=num_3_2, kernel=c(3, 3), pad=c(1, 1), 
             name=paste0(name, "_ra_3"), suffix="_conv_2")
  ra3 <- Conv(data=ra3, num_filter=num_3_3, kernel=c(3, 3), stride=c(2, 2), 
             name=paste0(name, "_ra_3"), suffix="_conv_3", withRelu=FALSE)
  
  merge_lst <- list()
  merge_lst <- c(ra1, ra2, ra3)
  merge_lst$num.args <- length(merge_lst)
  merge_lst$name <- paste0(name, "_ra_concat1")
  m <- mxnet:::mx.varg.symbol.Concat(merge_lst)
  
  m <- mx.symbol.BatchNorm(data=m, name=paste0(name, "_ra_bn1"))
  m <- mx.symbol.Activation(data=m, act_type="relu", name=paste0(name, "_ra_relu1"))
  
  return(m)
}

ReductionResnetB <- function(data,
                             num_2_1, num_2_2,
                             num_3_1, num_3_2,
                             num_4_1, num_4_2, num_4_3,
                             name){
  rb1 <- mx.symbol.Pooling(data=data, kernel=c(3, 3), stride=c(2, 2), 
                          pool_type="max", name=paste0("max_", name, "_pool1"))
  
  rb2 <- Conv(data=data, num_filter=num_2_1, name=paste0(name, "_rb_2"), suffix="_conv_1")
  rb2 <- Conv(data=rb2, num_filter=num_2_2, kernel=c(3, 3), stride=c(2, 2), 
             name=paste0(name, "_rb_2"), suffix="_conv_2", withRelu=FALSE)
  
  rb3 <- Conv(data=data, num_filter=num_3_1, name=paste0(name, "_rb_3"), suffix="_conv_1")
  rb3 <- Conv(data=rb3, num_filter=num_3_2, kernel=c(3, 3), stride=c(2, 2), 
             name=paste0(name, "_rb_3"), suffix="_conv_2", withRelu=FALSE)
  
  rb4 <- Conv(data=data, num_filter=num_4_1, name=paste0(name, "_rb_4"), suffix="_conv_1")
  rb4 <- Conv(data=rb4, num_filter=num_4_2, kernel=c(3, 3), pad=c(1, 1), 
             name=paste0(name, "_rb_4"), suffix="_conv_2")
  rb4 <- Conv(data=rb4, num_filter=num_4_3, kernel=c(3, 3), stride=c(2, 2), 
             name=paste0(name, "_rb_4"), suffix="_conv_3", withRelu=FALSE)
  
  merge_lst <- list()
  merge_lst <- c(rb1, rb2, rb3, rb4)
  merge_lst$num.args <- length(merge_lst)
  merge_lst$name <- paste0(name, "_rb_concat1")
  m <- mxnet:::mx.varg.symbol.Concat(merge_lst)
  
  m <- mx.symbol.BatchNorm(data=m, name=paste0(name, "_rb_bn1"))
  m <- mx.symbol.Activation(data=m, act_type="relu", name=paste0(name, "_rb_relu1"))
  
  return(m)
}

circle_in3a <- function(data,
                        num_1_1,
                        num_2_1, num_2_2,
                        num_3_1, num_3_2, num_3_3,
                        proj,
                        name,
                        scale,
                        round){
  in3a <- data
  for(i in 1:round){
    in3a <- InceptionResnetA(in3a,
                            num_1_1,
                            num_2_1, num_2_2,
                            num_3_1, num_3_2, num_3_3,
                            proj,
                            paste0(name, "_", i),
                            scaleResidual=scale)
  }
  return(in3a)
  
}

circle_in2b <- function(data,
                        num_1_1,
                        num_2_1, num_2_2, num_2_3,
                        proj,
                        name,
                        scale,
                        round){
  in2b <- data
  for(i in 1:round){
    in2b <- InceptionResnetB(in2b,
                            num_1_1,
                            num_2_1, num_2_2, num_2_3,
                            proj,
                            paste0(name, "_", i),
                            scaleResidual=scale)
  }
  return(in2b)
}

circle_in2c <- function(data,
                        num_1_1,
                        num_2_1, num_2_2, num_2_3,
                        proj,
                        name,
                        scale,
                        round){
  in2c <- data
  for(i in 1:round){
    in2c <- InceptionResnetC(in2c,
                            num_1_1,
                            num_2_1, num_2_2, num_2_3,
                            proj,
                            paste0(name, "_", i),
                            scaleResidual=scale)
  }
  return(in2c)
}

# create inception-resnet-v1
get_symbol <- function(num_classes=1000, scale=TRUE){
  
  # input shape 229*229*3
  data <- mx.symbol.Variable(name="data")
  
  # stage stem
  num_1_1 <- 32
  num_1_2 <- 32
  num_1_3 <- 64
  num_2_1 <- 80
  num_2_2 <- 192
  num_2_3 <- 256
  
  in_stem <- InceptionResnetStem(data,
                                num_1_1, num_1_2, num_1_3,
                                num_2_1, num_2_2, num_2_3,
                                "stem_stage")
  
  # stage 5 x Inception Resnet A
  num_1_1 <- 32
  num_2_1 <- 32
  num_2_2 <- 32
  num_3_1 <- 32
  num_3_2 <- 32
  num_3_3 <- 32
  proj <- 256
  
  in3a <- circle_in3a(in_stem,
                     num_1_1,
                     num_2_1, num_2_2,
                     num_3_1, num_3_2, num_3_3,
                     proj,
                     "in3a",
                     scale,
                     5)
  
  # stage Reduction Resnet A
  num_1_1 <- 384
  num_2_1 <- 192
  num_2_2 <- 192
  num_2_3 <- 256
  
  re3a <- ReductionResnetA(in3a,
                          num_1_1,
                          num_2_1, num_2_2, num_2_3,
                          "re3a")
  
  # stage 10 x Inception Resnet B
  num_1_1 <- 128
  num_2_1 <- 128
  num_2_2 <- 128
  num_2_3 <- 128
  proj <- 896
  
  in2b <- circle_in2b(re3a,
                     num_1_1,
                     num_2_1, num_2_2, num_2_3,
                     proj,
                     "in2b",
                     scale,
                     10)
  
  # stage Reduction Resnet B
  num_1_1 <- 256
  num_1_2 <- 384
  num_2_1 <- 256
  num_2_2 <- 256
  num_3_1 <- 256
  num_3_2 <- 256
  num_3_3 <- 256
  
  re4b <- ReductionResnetB(in2b,
                          num_1_1, num_1_2,
                          num_2_1, num_2_2,
                          num_3_1, num_3_2, num_3_3,
                          "re4b")
  
  # stage 5 x Inception Resnet C
  num_1_1 <- 128
  num_2_1 <- 192
  num_2_2 <- 192
  num_2_3 <- 192
  proj <-  1792
  
  in2c <- circle_in2c(re4b,
                     num_1_1,
                     num_2_1, num_2_2, num_2_3,
                     proj,
                     "in2c",
                     scale,
                     5)
  
  # stage Average Pooling
  pool <- mx.symbol.Pooling(data=in2c, kernel=c(8, 8), stride=c(1, 1), 
                           pool_type="avg", name="global_pool")
  
  # stage Dropout
  dropout <- mx.symbol.Dropout(data=pool, p=0.2)
  # dropout =  mx.symbol.Dropout(data=pool, p=0.8)
  flatten <- mx.symbol.Flatten(data=dropout, name="flatten")
  
  # output
  fc1 <- mx.symbol.FullyConnected(data=flatten, num_hidden=num_classes, name="fc1")
  softmax <- mx.symbol.SoftmaxOutput(data=fc1, name="softmax")
  
  return(softmax)
}