example-apps/PDD/pcb-defect-detection/libs/networks/slim_nets/nets_factory.py

   1 # Copyright 2016 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 """Contains a factory for building various models."""
  16
  17 from __future__ import absolute_import
  18 from __future__ import division
  19 from __future__ import print_function
  20 import functools
  21
  22 import tensorflow as tf
  23
  24 from nets import alexnet
  25 from nets import cifarnet
  26 from nets import inception
  27 from nets import lenet
  28 from nets import mobilenet_v1
  29 from nets import overfeat
  30 from nets import resnet_v1
  31 from nets import resnet_v2
  32 from nets import vgg
  33
  34 slim = tf.contrib.slim
  35
  36 networks_map = {'alexnet_v2': alexnet.alexnet_v2,
  37                 'cifarnet': cifarnet.cifarnet,
  38                 'overfeat': overfeat.overfeat,
  39                 'vgg_a': vgg.vgg_a,
  40                 'vgg_16': vgg.vgg_16,
  41                 'vgg_19': vgg.vgg_19,
  42                 'inception_v1': inception.inception_v1,
  43                 'inception_v2': inception.inception_v2,
  44                 'inception_v3': inception.inception_v3,
  45                 'inception_v4': inception.inception_v4,
  46                 'inception_resnet_v2': inception.inception_resnet_v2,
  47                 'lenet': lenet.lenet,
  48                 'resnet_v1_50': resnet_v1.resnet_v1_50,
  49                 'resnet_v1_101': resnet_v1.resnet_v1_101,
  50                 'resnet_v1_152': resnet_v1.resnet_v1_152,
  51                 'resnet_v1_200': resnet_v1.resnet_v1_200,
  52                 'resnet_v2_50': resnet_v2.resnet_v2_50,
  53                 'resnet_v2_101': resnet_v2.resnet_v2_101,
  54                 'resnet_v2_152': resnet_v2.resnet_v2_152,
  55                 'resnet_v2_200': resnet_v2.resnet_v2_200,
  56                 'mobilenet_v1': mobilenet_v1.mobilenet_v1,
  57                }
  58
  59 arg_scopes_map = {'alexnet_v2': alexnet.alexnet_v2_arg_scope,
  60                   'cifarnet': cifarnet.cifarnet_arg_scope,
  61                   'overfeat': overfeat.overfeat_arg_scope,
  62                   'vgg_a': vgg.vgg_arg_scope,
  63                   'vgg_16': vgg.vgg_arg_scope,
  64                   'vgg_19': vgg.vgg_arg_scope,
  65                   'inception_v1': inception.inception_v3_arg_scope,
  66                   'inception_v2': inception.inception_v3_arg_scope,
  67                   'inception_v3': inception.inception_v3_arg_scope,
  68                   'inception_v4': inception.inception_v4_arg_scope,
  69                   'inception_resnet_v2':
  70                   inception.inception_resnet_v2_arg_scope,
  71                   'lenet': lenet.lenet_arg_scope,
  72                   'resnet_v1_50': resnet_v1.resnet_arg_scope,
  73                   'resnet_v1_101': resnet_v1.resnet_arg_scope,
  74                   'resnet_v1_152': resnet_v1.resnet_arg_scope,
  75                   'resnet_v1_200': resnet_v1.resnet_arg_scope,
  76                   'resnet_v2_50': resnet_v2.resnet_arg_scope,
  77                   'resnet_v2_101': resnet_v2.resnet_arg_scope,
  78                   'resnet_v2_152': resnet_v2.resnet_arg_scope,
  79                   'resnet_v2_200': resnet_v2.resnet_arg_scope,
  80                   'mobilenet_v1': mobilenet_v1.mobilenet_v1_arg_scope,
  81                  }
  82
  83
  84 def get_network_fn(name, num_classes, weight_decay=0.0, is_training=False):
  85   """Returns a network_fn such as `logits, end_points = network_fn(images)`.
  86
  87   Args:
  88     name: The name of the network.
  89     num_classes: The number of classes to use for classification.
  90     weight_decay: The l2 coefficient for the model weights.
  91     is_training: `True` if the model is being used for training and `False`
  92       otherwise.
  93
  94   Returns:
  95     network_fn: A function that applies the model to a batch of images. It has
  96       the following signature:
  97         logits, end_points = network_fn(images)
  98   Raises:
  99     ValueError: If network `name` is not recognized.
 100   """
 101   if name not in networks_map:
 102     raise ValueError('Name of network unknown %s' % name)
 103   arg_scope = arg_scopes_map[name](weight_decay=weight_decay)
 104   func = networks_map[name]
 105   @functools.wraps(func)
 106   def network_fn(images):
 107     with slim.arg_scope(arg_scope):
 108       return func(images, num_classes, is_training=is_training)
 109   if hasattr(func, 'default_image_size'):
 110     network_fn.default_image_size = func.default_image_size
 111
 112   return network_fn