example-apps/PDD/pcb-defect-detection/libs/networks/slim_nets/alexnet.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 model definition for AlexNet.
  16
  17 This work was first described in:
  18   ImageNet Classification with Deep Convolutional Neural Networks
  19   Alex Krizhevsky, Ilya Sutskever and Geoffrey E. Hinton
  20
  21 and later refined in:
  22   One weird trick for parallelizing convolutional neural networks
  23   Alex Krizhevsky, 2014
  24
  25 Here we provide the implementation proposed in "One weird trick" and not
  26 "ImageNet Classification", as per the paper, the LRN layers have been removed.
  27
  28 Usage:
  29   with slim.arg_scope(alexnet.alexnet_v2_arg_scope()):
  30     outputs, end_points = alexnet.alexnet_v2(inputs)
  31
  32 @@alexnet_v2
  33 """
  34
  35 from __future__ import absolute_import
  36 from __future__ import division
  37 from __future__ import print_function
  38
  39 import tensorflow as tf
  40
  41 slim = tf.contrib.slim
  42 trunc_normal = lambda stddev: tf.truncated_normal_initializer(0.0, stddev)
  43
  44
  45 def alexnet_v2_arg_scope(weight_decay=0.0005):
  46   with slim.arg_scope([slim.conv2d, slim.fully_connected],
  47                       activation_fn=tf.nn.relu,
  48                       biases_initializer=tf.constant_initializer(0.1),
  49                       weights_regularizer=slim.l2_regularizer(weight_decay)):
  50     with slim.arg_scope([slim.conv2d], padding='SAME'):
  51       with slim.arg_scope([slim.max_pool2d], padding='VALID') as arg_sc:
  52         return arg_sc
  53
  54
  55 def alexnet_v2(inputs,
  56                num_classes=1000,
  57                is_training=True,
  58                dropout_keep_prob=0.5,
  59                spatial_squeeze=True,
  60                scope='alexnet_v2'):
  61   """AlexNet version 2.
  62
  63   Described in: http://arxiv.org/pdf/1404.5997v2.pdf
  64   Parameters from:
  65   github.com/akrizhevsky/cuda-convnet2/blob/master/layers/
  66   layers-imagenet-1gpu.cfg
  67
  68   Note: All the fully_connected layers have been transformed to conv2d layers.
  69         To use in classification mode, resize input to 224x224. To use in fully
  70         convolutional mode, set spatial_squeeze to false.
  71         The LRN layers have been removed and change the initializers from
  72         random_normal_initializer to xavier_initializer.
  73
  74   Args:
  75     inputs: a tensor of size [batch_size, height, width, channels].
  76     num_classes: number of predicted classes.
  77     is_training: whether or not the model is being trained.
  78     dropout_keep_prob: the probability that activations are kept in the dropout
  79       layers during training.
  80     spatial_squeeze: whether or not should squeeze the spatial dimensions of the
  81       outputs. Useful to remove unnecessary dimensions for classification.
  82     scope: Optional scope for the variables.
  83
  84   Returns:
  85     the last op containing the log predictions and end_points dict.
  86   """
  87   with tf.variable_scope(scope, 'alexnet_v2', [inputs]) as sc:
  88     end_points_collection = sc.name + '_end_points'
  89     # Collect outputs for conv2d, fully_connected and max_pool2d.
  90     with slim.arg_scope([slim.conv2d, slim.fully_connected, slim.max_pool2d],
  91                         outputs_collections=[end_points_collection]):
  92       net = slim.conv2d(inputs, 64, [11, 11], 4, padding='VALID',
  93                         scope='conv1')
  94       net = slim.max_pool2d(net, [3, 3], 2, scope='pool1')
  95       net = slim.conv2d(net, 192, [5, 5], scope='conv2')
  96       net = slim.max_pool2d(net, [3, 3], 2, scope='pool2')
  97       net = slim.conv2d(net, 384, [3, 3], scope='conv3')
  98       net = slim.conv2d(net, 384, [3, 3], scope='conv4')
  99       net = slim.conv2d(net, 256, [3, 3], scope='conv5')
 100       net = slim.max_pool2d(net, [3, 3], 2, scope='pool5')
 101
 102       # Use conv2d instead of fully_connected layers.
 103       with slim.arg_scope([slim.conv2d],
 104                           weights_initializer=trunc_normal(0.005),
 105                           biases_initializer=tf.constant_initializer(0.1)):
 106         net = slim.conv2d(net, 4096, [5, 5], padding='VALID',
 107                           scope='fc6')
 108         net = slim.dropout(net, dropout_keep_prob, is_training=is_training,
 109                            scope='dropout6')
 110         net = slim.conv2d(net, 4096, [1, 1], scope='fc7')
 111         net = slim.dropout(net, dropout_keep_prob, is_training=is_training,
 112                            scope='dropout7')
 113         net = slim.conv2d(net, num_classes, [1, 1],
 114                           activation_fn=None,
 115                           normalizer_fn=None,
 116                           biases_initializer=tf.zeros_initializer(),
 117                           scope='fc8')
 118
 119       # Convert end_points_collection into a end_point dict.
 120       end_points = slim.utils.convert_collection_to_dict(end_points_collection)
 121       if spatial_squeeze:
 122         net = tf.squeeze(net, [1, 2], name='fc8/squeezed')
 123         end_points[sc.name + '/fc8'] = net
 124       return net, end_points
 125 alexnet_v2.default_image_size = 224