example-apps/PDD/pcb-defect-detection/libs/networks/slim_nets/alexnet_test.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 """Tests for slim.slim_nets.alexnet."""
  16 from __future__ import absolute_import
  17 from __future__ import division
  18 from __future__ import print_function
  19
  20 import tensorflow as tf
  21
  22 from nets import alexnet
  23
  24 slim = tf.contrib.slim
  25
  26
  27 class AlexnetV2Test(tf.test.TestCase):
  28
  29   def testBuild(self):
  30     batch_size = 5
  31     height, width = 224, 224
  32     num_classes = 1000
  33     with self.test_session():
  34       inputs = tf.random_uniform((batch_size, height, width, 3))
  35       logits, _ = alexnet.alexnet_v2(inputs, num_classes)
  36       self.assertEquals(logits.op.name, 'alexnet_v2/fc8/squeezed')
  37       self.assertListEqual(logits.get_shape().as_list(),
  38                            [batch_size, num_classes])
  39
  40   def testFullyConvolutional(self):
  41     batch_size = 1
  42     height, width = 300, 400
  43     num_classes = 1000
  44     with self.test_session():
  45       inputs = tf.random_uniform((batch_size, height, width, 3))
  46       logits, _ = alexnet.alexnet_v2(inputs, num_classes, spatial_squeeze=False)
  47       self.assertEquals(logits.op.name, 'alexnet_v2/fc8/BiasAdd')
  48       self.assertListEqual(logits.get_shape().as_list(),
  49                            [batch_size, 4, 7, num_classes])
  50
  51   def testEndPoints(self):
  52     batch_size = 5
  53     height, width = 224, 224
  54     num_classes = 1000
  55     with self.test_session():
  56       inputs = tf.random_uniform((batch_size, height, width, 3))
  57       _, end_points = alexnet.alexnet_v2(inputs, num_classes)
  58       expected_names = ['alexnet_v2/conv1',
  59                         'alexnet_v2/pool1',
  60                         'alexnet_v2/conv2',
  61                         'alexnet_v2/pool2',
  62                         'alexnet_v2/conv3',
  63                         'alexnet_v2/conv4',
  64                         'alexnet_v2/conv5',
  65                         'alexnet_v2/pool5',
  66                         'alexnet_v2/fc6',
  67                         'alexnet_v2/fc7',
  68                         'alexnet_v2/fc8'
  69                        ]
  70       self.assertSetEqual(set(end_points.keys()), set(expected_names))
  71
  72   def testModelVariables(self):
  73     batch_size = 5
  74     height, width = 224, 224
  75     num_classes = 1000
  76     with self.test_session():
  77       inputs = tf.random_uniform((batch_size, height, width, 3))
  78       alexnet.alexnet_v2(inputs, num_classes)
  79       expected_names = ['alexnet_v2/conv1/weights',
  80                         'alexnet_v2/conv1/biases',
  81                         'alexnet_v2/conv2/weights',
  82                         'alexnet_v2/conv2/biases',
  83                         'alexnet_v2/conv3/weights',
  84                         'alexnet_v2/conv3/biases',
  85                         'alexnet_v2/conv4/weights',
  86                         'alexnet_v2/conv4/biases',
  87                         'alexnet_v2/conv5/weights',
  88                         'alexnet_v2/conv5/biases',
  89                         'alexnet_v2/fc6/weights',
  90                         'alexnet_v2/fc6/biases',
  91                         'alexnet_v2/fc7/weights',
  92                         'alexnet_v2/fc7/biases',
  93                         'alexnet_v2/fc8/weights',
  94                         'alexnet_v2/fc8/biases',
  95                        ]
  96       model_variables = [v.op.name for v in slim.get_model_variables()]
  97       self.assertSetEqual(set(model_variables), set(expected_names))
  98
  99   def testEvaluation(self):
 100     batch_size = 2
 101     height, width = 224, 224
 102     num_classes = 1000
 103     with self.test_session():
 104       eval_inputs = tf.random_uniform((batch_size, height, width, 3))
 105       logits, _ = alexnet.alexnet_v2(eval_inputs, is_training=False)
 106       self.assertListEqual(logits.get_shape().as_list(),
 107                            [batch_size, num_classes])
 108       predictions = tf.argmax(logits, 1)
 109       self.assertListEqual(predictions.get_shape().as_list(), [batch_size])
 110
 111   def testTrainEvalWithReuse(self):
 112     train_batch_size = 2
 113     eval_batch_size = 1
 114     train_height, train_width = 224, 224
 115     eval_height, eval_width = 300, 400
 116     num_classes = 1000
 117     with self.test_session():
 118       train_inputs = tf.random_uniform(
 119           (train_batch_size, train_height, train_width, 3))
 120       logits, _ = alexnet.alexnet_v2(train_inputs)
 121       self.assertListEqual(logits.get_shape().as_list(),
 122                            [train_batch_size, num_classes])
 123       tf.get_variable_scope().reuse_variables()
 124       eval_inputs = tf.random_uniform(
 125           (eval_batch_size, eval_height, eval_width, 3))
 126       logits, _ = alexnet.alexnet_v2(eval_inputs, is_training=False,
 127                                      spatial_squeeze=False)
 128       self.assertListEqual(logits.get_shape().as_list(),
 129                            [eval_batch_size, 4, 7, num_classes])
 130       logits = tf.reduce_mean(logits, [1, 2])
 131       predictions = tf.argmax(logits, 1)
 132       self.assertEquals(predictions.get_shape().as_list(), [eval_batch_size])
 133
 134   def testForward(self):
 135     batch_size = 1
 136     height, width = 224, 224
 137     with self.test_session() as sess:
 138       inputs = tf.random_uniform((batch_size, height, width, 3))
 139       logits, _ = alexnet.alexnet_v2(inputs)
 140       sess.run(tf.global_variables_initializer())
 141       output = sess.run(logits)
 142       self.assertTrue(output.any())
 143
 144 if __name__ == '__main__':
 145   tf.test.main()