example-apps/PDD/pcb-defect-detection/libs/networks/mobilenet/mobilenet_v2.py

   1 # Copyright 2018 The TensorFlow Authors. All Rights Reserved.
   2 #
   3 # Licensed under the Apache License, Version 2.0 (the "License");
   4 # you may not use this file except in compliance with the License.
   5 # You may obtain a copy of the License at
   6 #
   7 #     http://www.apache.org/licenses/LICENSE-2.0
   8 #
   9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS,
  11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12 # See the License for the specific language governing permissions and
  13 # limitations under the License.
  14 # ==============================================================================
  15 """Implementation of Mobilenet V2.
  16
  17 Architecture: https://arxiv.org/abs/1801.04381
  18
  19 The base model gives 72.2% accuracy on ImageNet, with 300MMadds,
  20 3.4 M parameters.
  21 """
  22
  23 from __future__ import absolute_import
  24 from __future__ import division
  25 from __future__ import print_function
  26
  27 import copy
  28
  29 import tensorflow as tf
  30
  31 from libs.networks.mobilenet import conv_blocks as ops
  32 from libs.networks.mobilenet import mobilenet as lib
  33
  34 slim = tf.contrib.slim
  35 op = lib.op
  36
  37 expand_input = ops.expand_input_by_factor
  38
  39 # pyformat: disable
  40 # Architecture: https://arxiv.org/abs/1801.04381
  41 V2_DEF = dict(
  42     defaults={
  43         # Note: these parameters of batch norm affect the architecture
  44         # that's why they are here and not in training_scope.
  45         (slim.batch_norm,): {'center': True, 'scale': True},
  46         (slim.conv2d, slim.fully_connected, slim.separable_conv2d): {
  47             'normalizer_fn': slim.batch_norm, 'activation_fn': tf.nn.relu6
  48         },
  49         (ops.expanded_conv,): {
  50             'expansion_size': expand_input(6),
  51             'split_expansion': 1,
  52             'normalizer_fn': slim.batch_norm,
  53             'residual': True
  54         },
  55         (slim.conv2d, slim.separable_conv2d): {'padding': 'SAME'}
  56     },
  57     spec=[
  58         op(slim.conv2d, stride=2, num_outputs=32, kernel_size=[3, 3]),
  59         op(ops.expanded_conv,
  60            expansion_size=expand_input(1, divisible_by=1),
  61            num_outputs=16),
  62         op(ops.expanded_conv, stride=2, num_outputs=24),
  63         op(ops.expanded_conv, stride=1, num_outputs=24),
  64         op(ops.expanded_conv, stride=2, num_outputs=32),
  65         op(ops.expanded_conv, stride=1, num_outputs=32),
  66         op(ops.expanded_conv, stride=1, num_outputs=32),
  67         op(ops.expanded_conv, stride=2, num_outputs=64),
  68         op(ops.expanded_conv, stride=1, num_outputs=64),
  69         op(ops.expanded_conv, stride=1, num_outputs=64),
  70         op(ops.expanded_conv, stride=1, num_outputs=64),
  71         op(ops.expanded_conv, stride=1, num_outputs=96),
  72         op(ops.expanded_conv, stride=1, num_outputs=96),
  73         op(ops.expanded_conv, stride=1, num_outputs=96),
  74         op(ops.expanded_conv, stride=2, num_outputs=160),
  75         op(ops.expanded_conv, stride=1, num_outputs=160),
  76         op(ops.expanded_conv, stride=1, num_outputs=160),
  77         op(ops.expanded_conv, stride=1, num_outputs=320),
  78         op(slim.conv2d, stride=1, kernel_size=[1, 1], num_outputs=1280)
  79     ],
  80 )
  81 # pyformat: enable
  82
  83
  84 @slim.add_arg_scope
  85 def mobilenet(input_tensor,
  86               num_classes=1001,
  87               depth_multiplier=1.0,
  88               scope='MobilenetV2',
  89               conv_defs=None,
  90               finegrain_classification_mode=False,
  91               min_depth=None,
  92               divisible_by=None,
  93               **kwargs):
  94   """Creates mobilenet V2 network.
  95
  96   Inference mode is created by default. To create training use training_scope
  97   below.
  98
  99   with tf.contrib.slim.arg_scope(mobilenet_v2.training_scope()):
 100      logits, endpoints = mobilenet_v2.mobilenet(input_tensor)
 101
 102   Args:
 103     input_tensor: The input tensor
 104     num_classes: number of classes
 105     depth_multiplier: The multiplier applied to scale number of
 106     channels in each layer. Note: this is called depth multiplier in the
 107     paper but the name is kept for consistency with slim's model builder.
 108     scope: Scope of the operator
 109     conv_defs: Allows to override default conv def.
 110     finegrain_classification_mode: When set to True, the model
 111     will keep the last layer large even for small multipliers. Following
 112     https://arxiv.org/abs/1801.04381
 113     suggests that it improves performance for ImageNet-type of problems.
 114       *Note* ignored if final_endpoint makes the builder exit earlier.
 115     min_depth: If provided, will ensure that all layers will have that
 116     many channels after application of depth multiplier.
 117     divisible_by: If provided will ensure that all layers # channels
 118     will be divisible by this number.
 119     **kwargs: passed directly to mobilenet.mobilenet:
 120       prediciton_fn- what prediction function to use.
 121       reuse-: whether to reuse variables (if reuse set to true, scope
 122       must be given).
 123   Returns:
 124     logits/endpoints pair
 125
 126   Raises:
 127     ValueError: On invalid arguments
 128   """
 129   if conv_defs is None:
 130     conv_defs = V2_DEF
 131   if 'multiplier' in kwargs:
 132     raise ValueError('mobilenetv2 doesn\'t support generic '
 133                      'multiplier parameter use "depth_multiplier" instead.')
 134   if finegrain_classification_mode:
 135     conv_defs = copy.deepcopy(conv_defs)
 136     if depth_multiplier < 1:
 137       conv_defs['spec'][-1].params['num_outputs'] /= depth_multiplier
 138
 139   depth_args = {}
 140   # NB: do not set depth_args unless they are provided to avoid overriding
 141   # whatever default depth_multiplier might have thanks to arg_scope.
 142   if min_depth is not None:
 143     depth_args['min_depth'] = min_depth
 144   if divisible_by is not None:
 145     depth_args['divisible_by'] = divisible_by
 146
 147   with slim.arg_scope((lib.depth_multiplier,), **depth_args):
 148     return lib.mobilenet(
 149         input_tensor,
 150         num_classes=num_classes,
 151         conv_defs=conv_defs,
 152         scope=scope,
 153         multiplier=depth_multiplier,
 154         **kwargs)
 155
 156
 157 @slim.add_arg_scope
 158 def mobilenet_base(input_tensor, depth_multiplier=1.0, **kwargs):
 159   """Creates base of the mobilenet (no pooling and no logits) ."""
 160   return mobilenet(input_tensor,
 161                    depth_multiplier=depth_multiplier,
 162                    base_only=True, **kwargs)
 163
 164
 165 def training_scope(**kwargs):
 166   """Defines MobilenetV2 training scope.
 167
 168   Usage:
 169      with tf.contrib.slim.arg_scope(mobilenet_v2.training_scope()):
 170        logits, endpoints = mobilenet_v2.mobilenet(input_tensor)
 171
 172   with slim.
 173
 174   Args:
 175     **kwargs: Passed to mobilenet.training_scope. The following parameters
 176     are supported:
 177       weight_decay- The weight decay to use for regularizing the model.
 178       stddev-  Standard deviation for initialization, if negative uses xavier.
 179       dropout_keep_prob- dropout keep probability
 180       bn_decay- decay for the batch norm moving averages.
 181
 182   Returns:
 183     An `arg_scope` to use for the mobilenet v2 model.
 184   """
 185   return lib.training_scope(**kwargs)
 186
 187
 188 __all__ = ['training_scope', 'mobilenet_base', 'mobilenet', 'V2_DEF']