def cifar10_resnet_v2_generator()

in sagemaker/byoc/tensorflow_bring_your_own_kr/container/cifar10/resnet_model.py [0:0]

• 37 lines of code
• 5 McCabe index (conditional complexity)

def cifar10_resnet_v2_generator(resnet_size, num_classes, data_format=None):
    """Generator for CIFAR-10 ResNet v2 models.

  Args:
    resnet_size: A single integer for the size of the ResNet model.
    num_classes: The number of possible classes for image classification.
    data_format: The input format ('channels_last', 'channels_first', or None).
      If set to None, the format is dependent on whether a GPU is available.

  Returns:
    The model function that takes in `inputs` and `is_training` and
    returns the output tensor of the ResNet model.

  Raises:
    ValueError: If `resnet_size` is invalid.
  """
    if resnet_size % 6 != 2:
        raise ValueError('resnet_size must be 6n + 2:', resnet_size)

    num_blocks = (resnet_size - 2) // 6

    if data_format is None:
        data_format = (
            'channels_first' if tf.test.is_built_with_cuda() else 'channels_last')

    def model(inputs, is_training):
        """Constructs the ResNet model given the inputs."""
        if data_format == 'channels_first':
            # Convert from channels_last (NHWC) to channels_first (NCHW). This
            # provides a large performance boost on GPU. See
            # https://www.tensorflow.org/performance/performance_guide#data_formats
            inputs = tf.transpose(inputs, [0, 3, 1, 2])

        inputs = conv2d_fixed_padding(
            inputs=inputs, filters=16, kernel_size=3, strides=1,
            data_format=data_format)
        inputs = tf.identity(inputs, 'initial_conv')

        inputs = block_layer(
            inputs=inputs, filters=16, block_fn=building_block, blocks=num_blocks,
            strides=1, is_training=is_training, name='block_layer1',
            data_format=data_format)
        inputs = block_layer(
            inputs=inputs, filters=32, block_fn=building_block, blocks=num_blocks,
            strides=2, is_training=is_training, name='block_layer2',
            data_format=data_format)
        inputs = block_layer(
            inputs=inputs, filters=64, block_fn=building_block, blocks=num_blocks,
            strides=2, is_training=is_training, name='block_layer3',
            data_format=data_format)

        inputs = batch_norm_relu(inputs, is_training, data_format)
        inputs = tf.layers.average_pooling2d(
            inputs=inputs, pool_size=8, strides=1, padding='VALID',
            data_format=data_format)
        inputs = tf.identity(inputs, 'final_avg_pool')
        inputs = tf.reshape(inputs, [-1, 64])
        inputs = tf.layers.dense(inputs=inputs, units=num_classes)
        inputs = tf.identity(inputs, 'final_dense')
        return inputs

    model.default_image_size = 32
    return model