Merge remote-tracking branch 'upstream/master'

research/audioset/vggish_train_demo.py

@@ -64,7 +64,7 @@ flags.DEFINE_integer(
 
 flags.DEFINE_boolean(
     'train_vggish', True,
-    'If Frue, allow VGGish parameters to change during training, thus '
+    'If True, allow VGGish parameters to change during training, thus '
     'fine-tuning VGGish. If False, VGGish parameters are fixed, thus using '
     'VGGish as a fixed feature extractor.')
 

research/autoaugment/README.md

+<font size=4><b>Train Wide-ResNet, Shake-Shake and ShakeDrop models on CIFAR-10
+and CIFAR-100 dataset with AutoAugment.</b></font>
+
+The CIFAR-10/CIFAR-100 data can be downloaded from:
+https://www.cs.toronto.edu/~kriz/cifar.html.
+
+The code replicates the results from Tables 1 and 2 on CIFAR-10/100 with the
+following models: Wide-ResNet-28-10, Shake-Shake (26 2x32d), Shake-Shake (26
+2x96d) and PyramidNet+ShakeDrop.
+
+<b>Related papers:</b>
+
+AutoAugment: Learning Augmentation Policies from Data
+
+https://arxiv.org/abs/1805.09501
+
+Wide Residual Networks
+
+https://arxiv.org/abs/1605.07146
+
+Shake-Shake regularization
+
+https://arxiv.org/abs/1705.07485
+
+ShakeDrop regularization
+
+https://arxiv.org/abs/1802.02375
+
+<b>Settings:</b>
+
+CIFAR-10 Model         | Learning Rate | Weight Decay | Num. Epochs | Batch Size
+---------------------- | ------------- | ------------ | ----------- | ----------
+Wide-ResNet-28-10      | 0.1           | 5e-4         | 200         | 128
+Shake-Shake (26 2x32d) | 0.01          | 1e-3         | 1800        | 128
+Shake-Shake (26 2x96d) | 0.01          | 1e-3         | 1800        | 128
+PyramidNet + ShakeDrop | 0.05          | 5e-5         | 1800        | 64
+
+<b>Prerequisite:</b>
+
+1.  Install TensorFlow.
+
+2.  Download CIFAR-10/CIFAR-100 dataset.
+
+```shell
+curl -o cifar-10-binary.tar.gz https://www.cs.toronto.edu/~kriz/cifar-10-binary.tar.gz
+curl -o cifar-100-binary.tar.gz https://www.cs.toronto.edu/~kriz/cifar-100-binary.tar.gz
+```
+
+<b>How to run:</b>
+
+```shell
+# cd to the your workspace.
+# Specify the directory where dataset is located using the data_path flag.
+# Note: User can split samples from training set into the eval set by changing train_size and validation_size.
+
+# For example, to train the Wide-ResNet-28-10 model on a GPU.
+python train_cifar.py --model_name=wrn \
+                      --checkpoint_dir=/tmp/training \
+                      --data_path=/tmp/data \
+                      --dataset='cifar10' \
+                      --use_cpu=0
+```
+
+## Contact for Issues
+
+*   Barret Zoph, @barretzoph <barretzoph@google.com>
+*   Ekin Dogus Cubuk, <cubuk@google.com>

research/autoaugment/augmentation_transforms.py

+# Copyright 2018 The TensorFlow Authors All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Transforms used in the Augmentation Policies."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import random
+import numpy as np
+# pylint:disable=g-multiple-import
+from PIL import ImageOps, ImageEnhance, ImageFilter, Image
+# pylint:enable=g-multiple-import
+
+
+IMAGE_SIZE = 32
+# What is the dataset mean and std of the images on the training set
+MEANS = [0.49139968, 0.48215841, 0.44653091]
+STDS = [0.24703223, 0.24348513, 0.26158784]
+PARAMETER_MAX = 10  # What is the max 'level' a transform could be predicted
+
+
+def random_flip(x):
+  """Flip the input x horizontally with 50% probability."""
+  if np.random.rand(1)[0] > 0.5:
+    return np.fliplr(x)
+  return x
+
+
+def zero_pad_and_crop(img, amount=4):
+  """Zero pad by `amount` zero pixels on each side then take a random crop.
+
+  Args:
+    img: numpy image that will be zero padded and cropped.
+    amount: amount of zeros to pad `img` with horizontally and verically.
+
+  Returns:
+    The cropped zero padded img. The returned numpy array will be of the same
+    shape as `img`.
+  """
+  padded_img = np.zeros((img.shape[0] + amount * 2, img.shape[1] + amount * 2,
+                         img.shape[2]))
+  padded_img[amount:img.shape[0] + amount, amount:
+             img.shape[1] + amount, :] = img
+  top = np.random.randint(low=0, high=2 * amount)
+  left = np.random.randint(low=0, high=2 * amount)
+  new_img = padded_img[top:top + img.shape[0], left:left + img.shape[1], :]
+  return new_img
+
+
+def create_cutout_mask(img_height, img_width, num_channels, size):
+  """Creates a zero mask used for cutout of shape `img_height` x `img_width`.
+
+  Args:
+    img_height: Height of image cutout mask will be applied to.
+    img_width: Width of image cutout mask will be applied to.
+    num_channels: Number of channels in the image.
+    size: Size of the zeros mask.
+
+  Returns:
+    A mask of shape `img_height` x `img_width` with all ones except for a
+    square of zeros of shape `size` x `size`. This mask is meant to be
+    elementwise multiplied with the original image. Additionally returns
+    the `upper_coord` and `lower_coord` which specify where the cutout mask
+    will be applied.
+  """
+  assert img_height == img_width
+
+  # Sample center where cutout mask will be applied
+  height_loc = np.random.randint(low=0, high=img_height)
+  width_loc = np.random.randint(low=0, high=img_width)
+
+  # Determine upper right and lower left corners of patch
+  upper_coord = (max(0, height_loc - size // 2), max(0, width_loc - size // 2))
+  lower_coord = (min(img_height, height_loc + size // 2),
+                 min(img_width, width_loc + size // 2))
+  mask_height = lower_coord[0] - upper_coord[0]
+  mask_width = lower_coord[1] - upper_coord[1]
+  assert mask_height > 0
+  assert mask_width > 0
+
+  mask = np.ones((img_height, img_width, num_channels))
+  zeros = np.zeros((mask_height, mask_width, num_channels))
+  mask[upper_coord[0]:lower_coord[0], upper_coord[1]:lower_coord[1], :] = (
+      zeros)
+  return mask, upper_coord, lower_coord
+
+
+def cutout_numpy(img, size=16):
+  """Apply cutout with mask of shape `size` x `size` to `img`.
+
+  The cutout operation is from the paper https://arxiv.org/abs/1708.04552.
+  This operation applies a `size`x`size` mask of zeros to a random location
+  within `img`.
+
+  Args:
+    img: Numpy image that cutout will be applied to.
+    size: Height/width of the cutout mask that will be
+
+  Returns:
+    A numpy tensor that is the result of applying the cutout mask to `img`.
+  """
+  img_height, img_width, num_channels = (img.shape[0], img.shape[1],
+                                         img.shape[2])
+  assert len(img.shape) == 3
+  mask, _, _ = create_cutout_mask(img_height, img_width, num_channels, size)
+  return img * mask
+
+
+def float_parameter(level, maxval):
+  """Helper function to scale `val` between 0 and maxval .
+
+  Args:
+    level: Level of the operation that will be between [0, `PARAMETER_MAX`].
+    maxval: Maximum value that the operation can have. This will be scaled
+      to level/PARAMETER_MAX.
+
+  Returns:
+    A float that results from scaling `maxval` according to `level`.
+  """
+  return float(level) * maxval / PARAMETER_MAX
+
+
+def int_parameter(level, maxval):
+  """Helper function to scale `val` between 0 and maxval .
+
+  Args:
+    level: Level of the operation that will be between [0, `PARAMETER_MAX`].
+    maxval: Maximum value that the operation can have. This will be scaled
+      to level/PARAMETER_MAX.
+
+  Returns:
+    An int that results from scaling `maxval` according to `level`.
+  """
+  return int(level * maxval / PARAMETER_MAX)
+
+
+def pil_wrap(img):
+  """Convert the `img` numpy tensor to a PIL Image."""
+  return Image.fromarray(
+      np.uint8((img * STDS + MEANS) * 255.0)).convert('RGBA')
+
+
+def pil_unwrap(pil_img):
+  """Converts the PIL img to a numpy array."""
+  pic_array = (np.array(pil_img.getdata()).reshape((32, 32, 4)) / 255.0)
+  i1, i2 = np.where(pic_array[:, :, 3] == 0)
+  pic_array = (pic_array[:, :, :3] - MEANS) / STDS
+  pic_array[i1, i2] = [0, 0, 0]
+  return pic_array
+
+
+def apply_policy(policy, img):
+  """Apply the `policy` to the numpy `img`.
+
+  Args:
+    policy: A list of tuples with the form (name, probability, level) where
+      `name` is the name of the augmentation operation to apply, `probability`
+      is the probability of applying the operation and `level` is what strength
+      the operation to apply.
+    img: Numpy image that will have `policy` applied to it.
+
+  Returns:
+    The result of applying `policy` to `img`.
+  """
+  pil_img = pil_wrap(img)
+  for xform in policy:
+    assert len(xform) == 3
+    name, probability, level = xform
+    xform_fn = NAME_TO_TRANSFORM[name].pil_transformer(probability, level)
+    pil_img = xform_fn(pil_img)
+  return pil_unwrap(pil_img)
+
+
+class TransformFunction(object):
+  """Wraps the Transform function for pretty printing options."""
+
+  def __init__(self, func, name):
+    self.f = func
+    self.name = name
+
+  def __repr__(self):
+    return '<' + self.name + '>'
+
+  def __call__(self, pil_img):
+    return self.f(pil_img)
+
+
+class TransformT(object):
+  """Each instance of this class represents a specific transform."""
+
+  def __init__(self, name, xform_fn):
+    self.name = name
+    self.xform = xform_fn
+
+  def pil_transformer(self, probability, level):
+
+    def return_function(im):
+      if random.random() < probability:
+        im = self.xform(im, level)
+      return im
+
+    name = self.name + '({:.1f},{})'.format(probability, level)
+    return TransformFunction(return_function, name)
+
+  def do_transform(self, image, level):
+    f = self.pil_transformer(PARAMETER_MAX, level)
+    return pil_unwrap(f(pil_wrap(image)))
+
+
+################## Transform Functions ##################
+identity = TransformT('identity', lambda pil_img, level: pil_img)
+flip_lr = TransformT(
+    'FlipLR',
+    lambda pil_img, level: pil_img.transpose(Image.FLIP_LEFT_RIGHT))
+flip_ud = TransformT(
+    'FlipUD',
+    lambda pil_img, level: pil_img.transpose(Image.FLIP_TOP_BOTTOM))
+# pylint:disable=g-long-lambda
+auto_contrast = TransformT(
+    'AutoContrast',
+    lambda pil_img, level: ImageOps.autocontrast(
+        pil_img.convert('RGB')).convert('RGBA'))
+equalize = TransformT(
+    'Equalize',
+    lambda pil_img, level: ImageOps.equalize(
+        pil_img.convert('RGB')).convert('RGBA'))
+invert = TransformT(
+    'Invert',
+    lambda pil_img, level: ImageOps.invert(
+        pil_img.convert('RGB')).convert('RGBA'))
+# pylint:enable=g-long-lambda
+blur = TransformT(
+    'Blur', lambda pil_img, level: pil_img.filter(ImageFilter.BLUR))
+smooth = TransformT(
+    'Smooth',
+    lambda pil_img, level: pil_img.filter(ImageFilter.SMOOTH))
+
+
+def _rotate_impl(pil_img, level):
+  """Rotates `pil_img` from -30 to 30 degrees depending on `level`."""
+  degrees = int_parameter(level, 30)
+  if random.random() > 0.5:
+    degrees = -degrees
+  return pil_img.rotate(degrees)
+
+
+rotate = TransformT('Rotate', _rotate_impl)
+
+
+def _posterize_impl(pil_img, level):
+  """Applies PIL Posterize to `pil_img`."""
+  level = int_parameter(level, 4)
+  return ImageOps.posterize(pil_img.convert('RGB'), 4 - level).convert('RGBA')
+
+
+posterize = TransformT('Posterize', _posterize_impl)
+
+
+def _shear_x_impl(pil_img, level):
+  """Applies PIL ShearX to `pil_img`.
+
+  The ShearX operation shears the image along the horizontal axis with `level`
+  magnitude.
+
+  Args:
+    pil_img: Image in PIL object.
+    level: Strength of the operation specified as an Integer from
+      [0, `PARAMETER_MAX`].
+
+  Returns:
+    A PIL Image that has had ShearX applied to it.
+  """
+  level = float_parameter(level, 0.3)
+  if random.random() > 0.5:
+    level = -level
+  return pil_img.transform((32, 32), Image.AFFINE, (1, level, 0, 0, 1, 0))
+
+
+shear_x = TransformT('ShearX', _shear_x_impl)
+
+
+def _shear_y_impl(pil_img, level):
+  """Applies PIL ShearY to `pil_img`.
+
+  The ShearY operation shears the image along the vertical axis with `level`
+  magnitude.
+
+  Args:
+    pil_img: Image in PIL object.
+    level: Strength of the operation specified as an Integer from
+      [0, `PARAMETER_MAX`].
+
+  Returns:
+    A PIL Image that has had ShearX applied to it.
+  """
+  level = float_parameter(level, 0.3)
+  if random.random() > 0.5:
+    level = -level
+  return pil_img.transform((32, 32), Image.AFFINE, (1, 0, 0, level, 1, 0))
+
+
+shear_y = TransformT('ShearY', _shear_y_impl)
+
+
+def _translate_x_impl(pil_img, level):
+  """Applies PIL TranslateX to `pil_img`.
+
+  Translate the image in the horizontal direction by `level`
+  number of pixels.
+
+  Args:
+    pil_img: Image in PIL object.
+    level: Strength of the operation specified as an Integer from
+      [0, `PARAMETER_MAX`].
+
+  Returns:
+    A PIL Image that has had TranslateX applied to it.
+  """
+  level = int_parameter(level, 10)
+  if random.random() > 0.5:
+    level = -level
+  return pil_img.transform((32, 32), Image.AFFINE, (1, 0, level, 0, 1, 0))
+
+
+translate_x = TransformT('TranslateX', _translate_x_impl)
+
+
+def _translate_y_impl(pil_img, level):
+  """Applies PIL TranslateY to `pil_img`.
+
+  Translate the image in the vertical direction by `level`
+  number of pixels.
+
+  Args:
+    pil_img: Image in PIL object.
+    level: Strength of the operation specified as an Integer from
+      [0, `PARAMETER_MAX`].
+
+  Returns:
+    A PIL Image that has had TranslateY applied to it.
+  """
+  level = int_parameter(level, 10)
+  if random.random() > 0.5:
+    level = -level
+  return pil_img.transform((32, 32), Image.AFFINE, (1, 0, 0, 0, 1, level))
+
+
+translate_y = TransformT('TranslateY', _translate_y_impl)
+
+
+def _crop_impl(pil_img, level, interpolation=Image.BILINEAR):
+  """Applies a crop to `pil_img` with the size depending on the `level`."""
+  cropped = pil_img.crop((level, level, IMAGE_SIZE - level, IMAGE_SIZE - level))
+  resized = cropped.resize((IMAGE_SIZE, IMAGE_SIZE), interpolation)
+  return resized
+
+
+crop_bilinear = TransformT('CropBilinear', _crop_impl)
+
+
+def _solarize_impl(pil_img, level):
+  """Applies PIL Solarize to `pil_img`.
+
+  Translate the image in the vertical direction by `level`
+  number of pixels.
+
+  Args:
+    pil_img: Image in PIL object.
+    level: Strength of the operation specified as an Integer from
+      [0, `PARAMETER_MAX`].
+
+  Returns:
+    A PIL Image that has had Solarize applied to it.
+  """
+  level = int_parameter(level, 256)
+  return ImageOps.solarize(pil_img.convert('RGB'), 256 - level).convert('RGBA')
+
+
+solarize = TransformT('Solarize', _solarize_impl)
+
+
+def _cutout_pil_impl(pil_img, level):
+  """Apply cutout to pil_img at the specified level."""
+  size = int_parameter(level, 20)
+  if size <= 0:
+    return pil_img
+  img_height, img_width, num_channels = (32, 32, 3)
+  _, upper_coord, lower_coord = (
+      create_cutout_mask(img_height, img_width, num_channels, size))
+  pixels = pil_img.load()  # create the pixel map
+  for i in range(upper_coord[0], lower_coord[0]):  # for every col:
+    for j in range(upper_coord[1], lower_coord[1]):  # For every row
+      pixels[i, j] = (125, 122, 113, 0)  # set the colour accordingly
+  return pil_img
+
+cutout = TransformT('Cutout', _cutout_pil_impl)
+
+
+def _enhancer_impl(enhancer):
+  """Sets level to be between 0.1 and 1.8 for ImageEnhance transforms of PIL."""
+  def impl(pil_img, level):
+    v = float_parameter(level, 1.8) + .1  # going to 0 just destroys it
+    return enhancer(pil_img).enhance(v)
+  return impl
+
+
+color = TransformT('Color', _enhancer_impl(ImageEnhance.Color))
+contrast = TransformT('Contrast', _enhancer_impl(ImageEnhance.Contrast))
+brightness = TransformT('Brightness', _enhancer_impl(
+    ImageEnhance.Brightness))
+sharpness = TransformT('Sharpness', _enhancer_impl(ImageEnhance.Sharpness))
+
+ALL_TRANSFORMS = [
+    flip_lr,
+    flip_ud,
+    auto_contrast,
+    equalize,
+    invert,
+    rotate,
+    posterize,
+    crop_bilinear,
+    solarize,
+    color,
+    contrast,
+    brightness,
+    sharpness,
+    shear_x,
+    shear_y,
+    translate_x,
+    translate_y,
+    cutout,
+    blur,
+    smooth
+]
+
+NAME_TO_TRANSFORM = {t.name: t for t in ALL_TRANSFORMS}
+TRANSFORM_NAMES = NAME_TO_TRANSFORM.keys()

research/autoaugment/custom_ops.py

+# Copyright 2018 The TensorFlow Authors All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Contains convenience wrappers for typical Neural Network TensorFlow layers.
+
+   Ops that have different behavior during training or eval have an is_training
+   parameter.
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+
+import numpy as np
+import tensorflow as tf
+
+
+arg_scope = tf.contrib.framework.arg_scope
+
+
+def variable(name, shape, dtype, initializer, trainable):
+  """Returns a TF variable with the passed in specifications."""
+  var = tf.get_variable(
+      name,
+      shape=shape,
+      dtype=dtype,
+      initializer=initializer,
+      trainable=trainable)
+  return var
+
+
+def global_avg_pool(x, scope=None):
+  """Average pools away spatial height and width dimension of 4D tensor."""
+  assert x.get_shape().ndims == 4
+  with tf.name_scope(scope, 'global_avg_pool', [x]):
+    kernel_size = (1, int(x.shape[1]), int(x.shape[2]), 1)
+    squeeze_dims = (1, 2)
+    result = tf.nn.avg_pool(
+        x,
+        ksize=kernel_size,
+        strides=(1, 1, 1, 1),
+        padding='VALID',
+        data_format='NHWC')
+    return tf.squeeze(result, squeeze_dims)
+
+
+def zero_pad(inputs, in_filter, out_filter):
+  """Zero pads `input` tensor to have `out_filter` number of filters."""
+  outputs = tf.pad(inputs, [[0, 0], [0, 0], [0, 0],
+                            [(out_filter - in_filter) // 2,
+                             (out_filter - in_filter) // 2]])
+  return outputs
+
+
+@tf.contrib.framework.add_arg_scope
+def batch_norm(inputs,
+               decay=0.999,
+               center=True,
+               scale=False,
+               epsilon=0.001,
+               is_training=True,
+               reuse=None,
+               scope=None):
+  """Small wrapper around tf.contrib.layers.batch_norm."""
+  return tf.contrib.layers.batch_norm(
+      inputs,
+      decay=decay,
+      center=center,
+      scale=scale,
+      epsilon=epsilon,
+      activation_fn=None,
+      param_initializers=None,
+      updates_collections=tf.GraphKeys.UPDATE_OPS,
+      is_training=is_training,
+      reuse=reuse,
+      trainable=True,
+      fused=True,
+      data_format='NHWC',
+      zero_debias_moving_mean=False,
+      scope=scope)
+
+
+def stride_arr(stride_h, stride_w):
+  return [1, stride_h, stride_w, 1]
+
+
+@tf.contrib.framework.add_arg_scope
+def conv2d(inputs,
+           num_filters_out,
+           kernel_size,
+           stride=1,
+           scope=None,
+           reuse=None):
+  """Adds a 2D convolution.
+
+  conv2d creates a variable called 'weights', representing the convolutional
+  kernel, that is convolved with the input.
+
+  Args:
+    inputs: a 4D tensor in NHWC format.
+    num_filters_out: the number of output filters.
+    kernel_size: an int specifying the kernel height and width size.
+    stride: an int specifying the height and width stride.
+    scope: Optional scope for variable_scope.
+    reuse: whether or not the layer and its variables should be reused.
+  Returns:
+    a tensor that is the result of a convolution being applied to `inputs`.
+  """
+  with tf.variable_scope(scope, 'Conv', [inputs], reuse=reuse):
+    num_filters_in = int(inputs.shape[3])
+    weights_shape = [kernel_size, kernel_size, num_filters_in, num_filters_out]
+
+    # Initialization
+    n = int(weights_shape[0] * weights_shape[1] * weights_shape[3])
+    weights_initializer = tf.random_normal_initializer(
+        stddev=np.sqrt(2.0 / n))
+
+    weights = variable(
+        name='weights',
+        shape=weights_shape,
+        dtype=tf.float32,
+        initializer=weights_initializer,
+        trainable=True)
+    strides = stride_arr(stride, stride)
+    outputs = tf.nn.conv2d(
+        inputs, weights, strides, padding='SAME', data_format='NHWC')
+    return outputs
+
+
+@tf.contrib.framework.add_arg_scope
+def fc(inputs,
+       num_units_out,
+       scope=None,
+       reuse=None):
+  """Creates a fully connected layer applied to `inputs`.
+
+  Args:
+    inputs: a tensor that the fully connected layer will be applied to. It
+      will be reshaped if it is not 2D.
+    num_units_out: the number of output units in the layer.
+    scope: Optional scope for variable_scope.
+    reuse: whether or not the layer and its variables should be reused.
+
+  Returns:
+     a tensor that is the result of applying a linear matrix to `inputs`.
+  """
+  if len(inputs.shape) > 2:
+    inputs = tf.reshape(inputs, [int(inputs.shape[0]), -1])
+
+  with tf.variable_scope(scope, 'FC', [inputs], reuse=reuse):
+    num_units_in = inputs.shape[1]
+    weights_shape = [num_units_in, num_units_out]
+    unif_init_range = 1.0 / (num_units_out)**(0.5)
+    weights_initializer = tf.random_uniform_initializer(
+        -unif_init_range, unif_init_range)
+    weights = variable(
+        name='weights',
+        shape=weights_shape,
+        dtype=tf.float32,
+        initializer=weights_initializer,
+        trainable=True)
+    bias_initializer = tf.constant_initializer(0.0)
+    biases = variable(
+        name='biases',
+        shape=[num_units_out,],
+        dtype=tf.float32,
+        initializer=bias_initializer,
+        trainable=True)
+    outputs = tf.nn.xw_plus_b(inputs, weights, biases)
+    return outputs
+
+
+@tf.contrib.framework.add_arg_scope
+def avg_pool(inputs, kernel_size, stride=2, padding='VALID', scope=None):
+  """Wrapper around tf.nn.avg_pool."""
+  with tf.name_scope(scope, 'AvgPool', [inputs]):
+    kernel = stride_arr(kernel_size, kernel_size)
+    strides = stride_arr(stride, stride)
+    return tf.nn.avg_pool(
+        inputs,
+        ksize=kernel,
+        strides=strides,
+        padding=padding,
+        data_format='NHWC')
+

research/autoaugment/data_utils.py

+# Copyright 2018 The TensorFlow Authors All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Data utils for CIFAR-10 and CIFAR-100."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import copy
+import cPickle
+import os
+import augmentation_transforms
+import numpy as np
+import policies as found_policies
+import tensorflow as tf
+
+
+# pylint:disable=logging-format-interpolation
+
+
+class DataSet(object):
+  """Dataset object that produces augmented training and eval data."""
+
+  def __init__(self, hparams):
+    self.hparams = hparams
+    self.epochs = 0
+    self.curr_train_index = 0
+
+    all_labels = []
+
+    self.good_policies = found_policies.good_policies()
+
+    # Determine how many databatched to load
+    num_data_batches_to_load = 5
+    total_batches_to_load = num_data_batches_to_load
+    train_batches_to_load = total_batches_to_load
+    assert hparams.train_size + hparams.validation_size <= 50000
+    if hparams.eval_test:
+      total_batches_to_load += 1
+    # Determine how many images we have loaded
+    total_dataset_size = 10000 * num_data_batches_to_load
+    train_dataset_size = total_dataset_size
+    if hparams.eval_test:
+      total_dataset_size += 10000
+
+    if hparams.dataset == 'cifar10':
+      all_data = np.empty((total_batches_to_load, 10000, 3072), dtype=np.uint8)
+    elif hparams.dataset == 'cifar100':
+      assert num_data_batches_to_load == 5
+      all_data = np.empty((1, 50000, 3072), dtype=np.uint8)
+      if hparams.eval_test:
+        test_data = np.empty((1, 10000, 3072), dtype=np.uint8)
+    if hparams.dataset == 'cifar10':
+      tf.logging.info('Cifar10')
+      datafiles = [
+          'data_batch_1', 'data_batch_2', 'data_batch_3', 'data_batch_4',
+          'data_batch_5']
+
+      datafiles = datafiles[:train_batches_to_load]
+      if hparams.eval_test:
+        datafiles.append('test_batch')
+      num_classes = 10
+    elif hparams.dataset == 'cifar100':
+      datafiles = ['train']
+      if hparams.eval_test:
+        datafiles.append('test')
+      num_classes = 100
+    else:
+      raise NotImplementedError('Unimplemented dataset: ', hparams.dataset)
+    if hparams.dataset != 'test':
+      for file_num, f in enumerate(datafiles):
+        d = unpickle(os.path.join(hparams.data_path, f))
+        if f == 'test':
+          test_data[0] = copy.deepcopy(d['data'])
+          all_data = np.concatenate([all_data, test_data], axis=1)
+        else:
+          all_data[file_num] = copy.deepcopy(d['data'])
+        if hparams.dataset == 'cifar10':
+          labels = np.array(d['labels'])
+        else:
+          labels = np.array(d['fine_labels'])
+        nsamples = len(labels)
+        for idx in range(nsamples):
+          all_labels.append(labels[idx])
+
+    all_data = all_data.reshape(total_dataset_size, 3072)
+    all_data = all_data.reshape(-1, 3, 32, 32)
+    all_data = all_data.transpose(0, 2, 3, 1).copy()
+    all_data = all_data / 255.0
+    mean = augmentation_transforms.MEANS
+    std = augmentation_transforms.STDS
+    tf.logging.info('mean:{}    std: {}'.format(mean, std))
+
+    all_data = (all_data - mean) / std
+    all_labels = np.eye(num_classes)[np.array(all_labels, dtype=np.int32)]
+    assert len(all_data) == len(all_labels)
+    tf.logging.info(
+        'In CIFAR10 loader, number of images: {}'.format(len(all_data)))
+
+    # Break off test data
+    if hparams.eval_test:
+      self.test_images = all_data[train_dataset_size:]
+      self.test_labels = all_labels[train_dataset_size:]
+
+    # Shuffle the rest of the data
+    all_data = all_data[:train_dataset_size]
+    all_labels = all_labels[:train_dataset_size]
+    np.random.seed(0)
+    perm = np.arange(len(all_data))
+    np.random.shuffle(perm)
+    all_data = all_data[perm]
+    all_labels = all_labels[perm]
+
+    # Break into train and val
+    train_size, val_size = hparams.train_size, hparams.validation_size
+    assert 50000 >= train_size + val_size
+    self.train_images = all_data[:train_size]
+    self.train_labels = all_labels[:train_size]
+    self.val_images = all_data[train_size:train_size + val_size]
+    self.val_labels = all_labels[train_size:train_size + val_size]
+    self.num_train = self.train_images.shape[0]
+
+  def next_batch(self):
+    """Return the next minibatch of augmented data."""
+    next_train_index = self.curr_train_index + self.hparams.batch_size
+    if next_train_index > self.num_train:
+      # Increase epoch number
+      epoch = self.epochs + 1
+      self.reset()
+      self.epochs = epoch
+    batched_data = (
+        self.train_images[self.curr_train_index:
+                          self.curr_train_index + self.hparams.batch_size],
+        self.train_labels[self.curr_train_index:
+                          self.curr_train_index + self.hparams.batch_size])
+    final_imgs = []
+
+    images, labels = batched_data
+    for data in images:
+      epoch_policy = self.good_policies[np.random.choice(
+          len(self.good_policies))]
+      final_img = augmentation_transforms.apply_policy(
+          epoch_policy, data)
+      final_img = augmentation_transforms.random_flip(
+          augmentation_transforms.zero_pad_and_crop(final_img, 4))
+      # Apply cutout
+      final_img = augmentation_transforms.cutout_numpy(final_img)
+      final_imgs.append(final_img)
+    batched_data = (np.array(final_imgs, np.float32), labels)
+    self.curr_train_index += self.hparams.batch_size
+    return batched_data
+
+  def reset(self):
+    """Reset training data and index into the training data."""
+    self.epochs = 0
+    # Shuffle the training data
+    perm = np.arange(self.num_train)
+    np.random.shuffle(perm)
+    assert self.num_train == self.train_images.shape[
+        0], 'Error incorrect shuffling mask'
+    self.train_images = self.train_images[perm]
+    self.train_labels = self.train_labels[perm]
+    self.curr_train_index = 0
+
+
+def unpickle(f):
+  tf.logging.info('loading file: {}'.format(f))
+  fo = tf.gfile.Open(f, 'r')
+  d = cPickle.load(fo)
+  fo.close()
+  return d

research/autoaugment/helper_utils.py

+# Copyright 2018 The TensorFlow Authors All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Helper functions used for training AutoAugment models."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+import tensorflow as tf
+
+
+def setup_loss(logits, labels):
+  """Returns the cross entropy for the given `logits` and `labels`."""
+  predictions = tf.nn.softmax(logits)
+  cost = tf.losses.softmax_cross_entropy(onehot_labels=labels,
+                                         logits=logits)
+  return predictions, cost
+
+
+def decay_weights(cost, weight_decay_rate):
+  """Calculates the loss for l2 weight decay and adds it to `cost`."""
+  costs = []
+  for var in tf.trainable_variables():
+    costs.append(tf.nn.l2_loss(var))
+  cost += tf.multiply(weight_decay_rate, tf.add_n(costs))
+  return cost
+
+
+def eval_child_model(session, model, data_loader, mode):
+  """Evaluates `model` on held out data depending on `mode`.
+
+  Args:
+    session: TensorFlow session the model will be run with.
+    model: TensorFlow model that will be evaluated.
+    data_loader: DataSet object that contains data that `model` will
+      evaluate.
+    mode: Will `model` either evaluate validation or test data.
+
+  Returns:
+    Accuracy of `model` when evaluated on the specified dataset.
+
+  Raises:
+    ValueError: if invalid dataset `mode` is specified.
+  """
+  if mode == 'val':
+    images = data_loader.val_images
+    labels = data_loader.val_labels
+  elif mode == 'test':
+    images = data_loader.test_images
+    labels = data_loader.test_labels
+  else:
+    raise ValueError('Not valid eval mode')
+  assert len(images) == len(labels)
+  tf.logging.info('model.batch_size is {}'.format(model.batch_size))
+  assert len(images) % model.batch_size == 0
+  eval_batches = int(len(images) / model.batch_size)
+  for i in range(eval_batches):
+    eval_images = images[i * model.batch_size:(i + 1) * model.batch_size]
+    eval_labels = labels[i * model.batch_size:(i + 1) * model.batch_size]
+    _ = session.run(
+        model.eval_op,
+        feed_dict={
+            model.images: eval_images,
+            model.labels: eval_labels,
+        })
+  return session.run(model.accuracy)
+
+
+def cosine_lr(learning_rate, epoch, iteration, batches_per_epoch, total_epochs):
+  """Cosine Learning rate.
+
+  Args:
+    learning_rate: Initial learning rate.
+    epoch: Current epoch we are one. This is one based.
+    iteration: Current batch in this epoch.
+    batches_per_epoch: Batches per epoch.
+    total_epochs: Total epochs you are training for.
+
+  Returns:
+    The learning rate to be used for this current batch.
+  """
+  t_total = total_epochs * batches_per_epoch
+  t_cur = float(epoch * batches_per_epoch + iteration)
+  return 0.5 * learning_rate * (1 + np.cos(np.pi * t_cur / t_total))
+
+
+def get_lr(curr_epoch, hparams, iteration=None):
+  """Returns the learning rate during training based on the current epoch."""
+  assert iteration is not None
+  batches_per_epoch = int(hparams.train_size / hparams.batch_size)
+  lr = cosine_lr(hparams.lr, curr_epoch, iteration, batches_per_epoch,
+                 hparams.num_epochs)
+  return lr
+
+
+def run_epoch_training(session, model, data_loader, curr_epoch):
+  """Runs one epoch of training for the model passed in.
+
+  Args:
+    session: TensorFlow session the model will be run with.
+    model: TensorFlow model that will be evaluated.
+    data_loader: DataSet object that contains data that `model` will
+      evaluate.
+    curr_epoch: How many of epochs of training have been done so far.
+
+  Returns:
+    The accuracy of 'model' on the training set
+  """
+  steps_per_epoch = int(model.hparams.train_size / model.hparams.batch_size)
+  tf.logging.info('steps per epoch: {}'.format(steps_per_epoch))
+  curr_step = session.run(model.global_step)
+  assert curr_step % steps_per_epoch == 0
+
+  # Get the current learning rate for the model based on the current epoch
+  curr_lr = get_lr(curr_epoch, model.hparams, iteration=0)
+  tf.logging.info('lr of {} for epoch {}'.format(curr_lr, curr_epoch))
+
+  for step in xrange(steps_per_epoch):
+    curr_lr = get_lr(curr_epoch, model.hparams, iteration=(step + 1))
+    # Update the lr rate variable to the current LR.
+    model.lr_rate_ph.load(curr_lr, session=session)
+    if step % 20 == 0:
+      tf.logging.info('Training {}/{}'.format(step, steps_per_epoch))
+
+    train_images, train_labels = data_loader.next_batch()
+    _, step, _ = session.run(
+        [model.train_op, model.global_step, model.eval_op],
+        feed_dict={
+            model.images: train_images,
+            model.labels: train_labels,
+        })
+
+  train_accuracy = session.run(model.accuracy)
+  tf.logging.info('Train accuracy: {}'.format(train_accuracy))
+  return train_accuracy

research/autoaugment/policies.py

+# Copyright 2018 The TensorFlow Authors All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+
+def good_policies():
+  """AutoAugment policies found on Cifar."""
+  exp0_0 = [
+      [('Invert', 0.1, 7), ('Contrast', 0.2, 6)],
+      [('Rotate', 0.7, 2), ('TranslateX', 0.3, 9)],
+      [('Sharpness', 0.8, 1), ('Sharpness', 0.9, 3)],
+      [('ShearY', 0.5, 8), ('TranslateY', 0.7, 9)],
+      [('AutoContrast', 0.5, 8), ('Equalize', 0.9, 2)]]
+  exp0_1 = [
+      [('Solarize', 0.4, 5), ('AutoContrast', 0.9, 3)],
+      [('TranslateY', 0.9, 9), ('TranslateY', 0.7, 9)],
+      [('AutoContrast', 0.9, 2), ('Solarize', 0.8, 3)],
+      [('Equalize', 0.8, 8), ('Invert', 0.1, 3)],
+      [('TranslateY', 0.7, 9), ('AutoContrast', 0.9, 1)]]
+  exp0_2 = [
+      [('Solarize', 0.4, 5), ('AutoContrast', 0.0, 2)],
+      [('TranslateY', 0.7, 9), ('TranslateY', 0.7, 9)],
+      [('AutoContrast', 0.9, 0), ('Solarize', 0.4, 3)],
+      [('Equalize', 0.7, 5), ('Invert', 0.1, 3)],
+      [('TranslateY', 0.7, 9), ('TranslateY', 0.7, 9)]]
+  exp0_3 = [
+      [('Solarize', 0.4, 5), ('AutoContrast', 0.9, 1)],
+      [('TranslateY', 0.8, 9), ('TranslateY', 0.9, 9)],
+      [('AutoContrast', 0.8, 0), ('TranslateY', 0.7, 9)],
+      [('TranslateY', 0.2, 7), ('Color', 0.9, 6)],
+      [('Equalize', 0.7, 6), ('Color', 0.4, 9)]]
+  exp1_0 = [
+      [('ShearY', 0.2, 7), ('Posterize', 0.3, 7)],
+      [('Color', 0.4, 3), ('Brightness', 0.6, 7)],
+      [('Sharpness', 0.3, 9), ('Brightness', 0.7, 9)],
+      [('Equalize', 0.6, 5), ('Equalize', 0.5, 1)],
+      [('Contrast', 0.6, 7), ('Sharpness', 0.6, 5)]]
+  exp1_1 = [
+      [('Brightness', 0.3, 7), ('AutoContrast', 0.5, 8)],
+      [('AutoContrast', 0.9, 4), ('AutoContrast', 0.5, 6)],
+      [('Solarize', 0.3, 5), ('Equalize', 0.6, 5)],
+      [('TranslateY', 0.2, 4), ('Sharpness', 0.3, 3)],
+      [('Brightness', 0.0, 8), ('Color', 0.8, 8)]]
+  exp1_2 = [
+      [('Solarize', 0.2, 6), ('Color', 0.8, 6)],
+      [('Solarize', 0.2, 6), ('AutoContrast', 0.8, 1)],
+      [('Solarize', 0.4, 1), ('Equalize', 0.6, 5)],
+      [('Brightness', 0.0, 0), ('Solarize', 0.5, 2)],
+      [('AutoContrast', 0.9, 5), ('Brightness', 0.5, 3)]]
+  exp1_3 = [
+      [('Contrast', 0.7, 5), ('Brightness', 0.0, 2)],
+      [('Solarize', 0.2, 8), ('Solarize', 0.1, 5)],
+      [('Contrast', 0.5, 1), ('TranslateY', 0.2, 9)],
+      [('AutoContrast', 0.6, 5), ('TranslateY', 0.0, 9)],
+      [('AutoContrast', 0.9, 4), ('Equalize', 0.8, 4)]]
+  exp1_4 = [
+      [('Brightness', 0.0, 7), ('Equalize', 0.4, 7)],
+      [('Solarize', 0.2, 5), ('Equalize', 0.7, 5)],
+      [('Equalize', 0.6, 8), ('Color', 0.6, 2)],
+      [('Color', 0.3, 7), ('Color', 0.2, 4)],
+      [('AutoContrast', 0.5, 2), ('Solarize', 0.7, 2)]]
+  exp1_5 = [
+      [('AutoContrast', 0.2, 0), ('Equalize', 0.1, 0)],
+      [('ShearY', 0.6, 5), ('Equalize', 0.6, 5)],
+      [('Brightness', 0.9, 3), ('AutoContrast', 0.4, 1)],
+      [('Equalize', 0.8, 8), ('Equalize', 0.7, 7)],
+      [('Equalize', 0.7, 7), ('Solarize', 0.5, 0)]]
+  exp1_6 = [
+      [('Equalize', 0.8, 4), ('TranslateY', 0.8, 9)],
+      [('TranslateY', 0.8, 9), ('TranslateY', 0.6, 9)],
+      [('TranslateY', 0.9, 0), ('TranslateY', 0.5, 9)],
+      [('AutoContrast', 0.5, 3), ('Solarize', 0.3, 4)],
+      [('Solarize', 0.5, 3), ('Equalize', 0.4, 4)]]
+  exp2_0 = [
+      [('Color', 0.7, 7), ('TranslateX', 0.5, 8)],
+      [('Equalize', 0.3, 7), ('AutoContrast', 0.4, 8)],
+      [('TranslateY', 0.4, 3), ('Sharpness', 0.2, 6)],
+      [('Brightness', 0.9, 6), ('Color', 0.2, 8)],
+      [('Solarize', 0.5, 2), ('Invert', 0.0, 3)]]
+  exp2_1 = [
+      [('AutoContrast', 0.1, 5), ('Brightness', 0.0, 0)],
+      [('Cutout', 0.2, 4), ('Equalize', 0.1, 1)],
+      [('Equalize', 0.7, 7), ('AutoContrast', 0.6, 4)],
+      [('Color', 0.1, 8), ('ShearY', 0.2, 3)],
+      [('ShearY', 0.4, 2), ('Rotate', 0.7, 0)]]
+  exp2_2 = [
+      [('ShearY', 0.1, 3), ('AutoContrast', 0.9, 5)],
+      [('TranslateY', 0.3, 6), ('Cutout', 0.3, 3)],
+      [('Equalize', 0.5, 0), ('Solarize', 0.6, 6)],
+      [('AutoContrast', 0.3, 5), ('Rotate', 0.2, 7)],
+      [('Equalize', 0.8, 2), ('Invert', 0.4, 0)]]
+  exp2_3 = [
+      [('Equalize', 0.9, 5), ('Color', 0.7, 0)],
+      [('Equalize', 0.1, 1), ('ShearY', 0.1, 3)],
+      [('AutoContrast', 0.7, 3), ('Equalize', 0.7, 0)],
+      [('Brightness', 0.5, 1), ('Contrast', 0.1, 7)],
+      [('Contrast', 0.1, 4), ('Solarize', 0.6, 5)]]
+  exp2_4 = [
+      [('Solarize', 0.2, 3), ('ShearX', 0.0, 0)],
+      [('TranslateX', 0.3, 0), ('TranslateX', 0.6, 0)],
+      [('Equalize', 0.5, 9), ('TranslateY', 0.6, 7)],
+      [('ShearX', 0.1, 0), ('Sharpness', 0.5, 1)],
+      [('Equalize', 0.8, 6), ('Invert', 0.3, 6)]]
+  exp2_5 = [
+      [('AutoContrast', 0.3, 9), ('Cutout', 0.5, 3)],
+      [('ShearX', 0.4, 4), ('AutoContrast', 0.9, 2)],
+      [('ShearX', 0.0, 3), ('Posterize', 0.0, 3)],
+      [('Solarize', 0.4, 3), ('Color', 0.2, 4)],
+      [('Equalize', 0.1, 4), ('Equalize', 0.7, 6)]]
+  exp2_6 = [
+      [('Equalize', 0.3, 8), ('AutoContrast', 0.4, 3)],
+      [('Solarize', 0.6, 4), ('AutoContrast', 0.7, 6)],
+      [('AutoContrast', 0.2, 9), ('Brightness', 0.4, 8)],
+      [('Equalize', 0.1, 0), ('Equalize', 0.0, 6)],
+      [('Equalize', 0.8, 4), ('Equalize', 0.0, 4)]]
+  exp2_7 = [
+      [('Equalize', 0.5, 5), ('AutoContrast', 0.1, 2)],
+      [('Solarize', 0.5, 5), ('AutoContrast', 0.9, 5)],
+      [('AutoContrast', 0.6, 1), ('AutoContrast', 0.7, 8)],
+      [('Equalize', 0.2, 0), ('AutoContrast', 0.1, 2)],
+      [('Equalize', 0.6, 9), ('Equalize', 0.4, 4)]]
+  exp0s = exp0_0 + exp0_1 + exp0_2 + exp0_3
+  exp1s = exp1_0 + exp1_1 + exp1_2 + exp1_3 + exp1_4 + exp1_5 + exp1_6
+  exp2s = exp2_0 + exp2_1 + exp2_2 + exp2_3 + exp2_4 + exp2_5 + exp2_6 + exp2_7
+  return  exp0s + exp1s + exp2s

research/autoaugment/shake_drop.py

+# Copyright 2018 The TensorFlow Authors All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Builds the Shake-Shake Model."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import math
+import custom_ops as ops
+import tensorflow as tf
+
+
+def round_int(x):
+  """Rounds `x` and then converts to an int."""
+  return int(math.floor(x + 0.5))
+
+
+def shortcut(x, output_filters, stride):
+  """Applies strided avg pool or zero padding to make output_filters match x."""
+  num_filters = int(x.shape[3])
+  if stride == 2:
+    x = ops.avg_pool(x, 2, stride=stride, padding='SAME')
+  if num_filters != output_filters:
+    diff = output_filters - num_filters
+    assert diff > 0
+    # Zero padd diff zeros
+    padding = [[0, 0], [0, 0], [0, 0], [0, diff]]
+    x = tf.pad(x, padding)
+  return x
+
+
+def calc_prob(curr_layer, total_layers, p_l):
+  """Calculates drop prob depending on the current layer."""
+  return 1 - (float(curr_layer) / total_layers) * p_l
+
+
+def bottleneck_layer(x, n, stride, prob, is_training, alpha, beta):
+  """Bottleneck layer for shake drop model."""
+  assert alpha[1] > alpha[0]
+  assert beta[1] > beta[0]
+  with tf.variable_scope('bottleneck_{}'.format(prob)):
+    input_layer = x
+    x = ops.batch_norm(x, scope='bn_1_pre')
+    x = ops.conv2d(x, n, 1, scope='1x1_conv_contract')
+    x = ops.batch_norm(x, scope='bn_1_post')
+    x = tf.nn.relu(x)
+    x = ops.conv2d(x, n, 3, stride=stride, scope='3x3')
+    x = ops.batch_norm(x, scope='bn_2')
+    x = tf.nn.relu(x)
+    x = ops.conv2d(x, n * 4, 1, scope='1x1_conv_expand')
+    x = ops.batch_norm(x, scope='bn_3')
+
+    # Apply regularization here
+    # Sample bernoulli with prob
+    if is_training:
+      batch_size = tf.shape(x)[0]
+      bern_shape = [batch_size, 1, 1, 1]
+      random_tensor = prob
+      random_tensor += tf.random_uniform(bern_shape, dtype=tf.float32)
+      binary_tensor = tf.floor(random_tensor)
+
+      alpha_values = tf.random_uniform(
+          [batch_size, 1, 1, 1], minval=alpha[0], maxval=alpha[1],
+          dtype=tf.float32)
+      beta_values = tf.random_uniform(
+          [batch_size, 1, 1, 1], minval=beta[0], maxval=beta[1],
+          dtype=tf.float32)
+      rand_forward = (
+          binary_tensor + alpha_values - binary_tensor * alpha_values)
+      rand_backward = (
+          binary_tensor + beta_values - binary_tensor * beta_values)
+      x = x * rand_backward + tf.stop_gradient(x * rand_forward -
+                                               x * rand_backward)
+    else:
+      expected_alpha = (alpha[1] + alpha[0])/2
+      # prob is the expectation of the bernoulli variable
+      x = (prob + expected_alpha - prob * expected_alpha) * x
+
+    res = shortcut(input_layer, n * 4, stride)
+    return x + res
+
+
+def build_shake_drop_model(images, num_classes, is_training):
+  """Builds the PyramidNet Shake-Drop model.
+
+  Build the PyramidNet Shake-Drop model from https://arxiv.org/abs/1802.02375.
+
+  Args:
+    images: Tensor of images that will be fed into the Wide ResNet Model.
+    num_classes: Number of classed that the model needs to predict.
+    is_training: Is the model training or not.
+
+  Returns:
+    The logits of the PyramidNet Shake-Drop model.
+  """
+  # ShakeDrop Hparams
+  p_l = 0.5
+  alpha_shake = [-1, 1]
+  beta_shake = [0, 1]
+
+  # PyramidNet Hparams
+  alpha = 200
+  depth = 272
+  # This is for the bottleneck architecture specifically
+  n = int((depth - 2) / 9)
+  start_channel = 16
+  add_channel = alpha / (3 * n)
+
+  # Building the models
+  x = images
+  x = ops.conv2d(x, 16, 3, scope='init_conv')
+  x = ops.batch_norm(x, scope='init_bn')
+
+  layer_num = 1
+  total_layers = n * 3
+  start_channel += add_channel
+  prob = calc_prob(layer_num, total_layers, p_l)
+  x = bottleneck_layer(
+      x, round_int(start_channel), 1, prob, is_training, alpha_shake,
+      beta_shake)
+  layer_num += 1
+  for _ in range(1, n):
+    start_channel += add_channel
+    prob = calc_prob(layer_num, total_layers, p_l)
+    x = bottleneck_layer(
+        x, round_int(start_channel), 1, prob, is_training, alpha_shake,
+        beta_shake)
+    layer_num += 1
+
+  start_channel += add_channel
+  prob = calc_prob(layer_num, total_layers, p_l)
+  x = bottleneck_layer(
+      x, round_int(start_channel), 2, prob, is_training, alpha_shake,
+      beta_shake)
+  layer_num += 1
+  for _ in range(1, n):
+    start_channel += add_channel
+    prob = calc_prob(layer_num, total_layers, p_l)
+    x = bottleneck_layer(
+        x, round_int(start_channel), 1, prob, is_training, alpha_shake,
+        beta_shake)
+    layer_num += 1
+
+  start_channel += add_channel
+  prob = calc_prob(layer_num, total_layers, p_l)
+  x = bottleneck_layer(
+      x, round_int(start_channel), 2, prob, is_training, alpha_shake,
+      beta_shake)
+  layer_num += 1
+  for _ in range(1, n):
+    start_channel += add_channel
+    prob = calc_prob(layer_num, total_layers, p_l)
+    x = bottleneck_layer(
+        x, round_int(start_channel), 1, prob, is_training, alpha_shake,
+        beta_shake)
+    layer_num += 1
+
+  assert layer_num - 1 == total_layers
+  x = ops.batch_norm(x, scope='final_bn')
+  x = tf.nn.relu(x)
+  x = ops.global_avg_pool(x)
+  # Fully connected
+  logits = ops.fc(x, num_classes)
+  return logits

research/autoaugment/shake_shake.py

+# Copyright 2018 The TensorFlow Authors All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Builds the Shake-Shake Model."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import custom_ops as ops
+import tensorflow as tf
+
+
+def _shake_shake_skip_connection(x, output_filters, stride):
+  """Adds a residual connection to the filter x for the shake-shake model."""
+  curr_filters = int(x.shape[3])
+  if curr_filters == output_filters:
+    return x
+  stride_spec = ops.stride_arr(stride, stride)
+  # Skip path 1
+  path1 = tf.nn.avg_pool(
+      x, [1, 1, 1, 1], stride_spec, 'VALID', data_format='NHWC')
+  path1 = ops.conv2d(path1, int(output_filters / 2), 1, scope='path1_conv')
+
+  # Skip path 2
+  # First pad with 0's then crop
+  pad_arr = [[0, 0], [0, 1], [0, 1], [0, 0]]
+  path2 = tf.pad(x, pad_arr)[:, 1:, 1:, :]
+  concat_axis = 3
+
+  path2 = tf.nn.avg_pool(
+      path2, [1, 1, 1, 1], stride_spec, 'VALID', data_format='NHWC')
+  path2 = ops.conv2d(path2, int(output_filters / 2), 1, scope='path2_conv')
+
+  # Concat and apply BN
+  final_path = tf.concat(values=[path1, path2], axis=concat_axis)
+  final_path = ops.batch_norm(final_path, scope='final_path_bn')
+  return final_path
+
+
+def _shake_shake_branch(x, output_filters, stride, rand_forward, rand_backward,
+                        is_training):
+  """Building a 2 branching convnet."""
+  x = tf.nn.relu(x)
+  x = ops.conv2d(x, output_filters, 3, stride=stride, scope='conv1')
+  x = ops.batch_norm(x, scope='bn1')
+  x = tf.nn.relu(x)
+  x = ops.conv2d(x, output_filters, 3, scope='conv2')
+  x = ops.batch_norm(x, scope='bn2')
+  if is_training:
+    x = x * rand_backward + tf.stop_gradient(x * rand_forward -
+                                             x * rand_backward)
+  else:
+    x *= 1.0 / 2
+  return x
+
+
+def _shake_shake_block(x, output_filters, stride, is_training):
+  """Builds a full shake-shake sub layer."""
+  batch_size = tf.shape(x)[0]
+
+  # Generate random numbers for scaling the branches
+  rand_forward = [
+      tf.random_uniform(
+          [batch_size, 1, 1, 1], minval=0, maxval=1, dtype=tf.float32)
+      for _ in range(2)
+  ]
+  rand_backward = [
+      tf.random_uniform(
+          [batch_size, 1, 1, 1], minval=0, maxval=1, dtype=tf.float32)
+      for _ in range(2)
+  ]
+  # Normalize so that all sum to 1
+  total_forward = tf.add_n(rand_forward)
+  total_backward = tf.add_n(rand_backward)
+  rand_forward = [samp / total_forward for samp in rand_forward]
+  rand_backward = [samp / total_backward for samp in rand_backward]
+  zipped_rand = zip(rand_forward, rand_backward)
+
+  branches = []
+  for branch, (r_forward, r_backward) in enumerate(zipped_rand):
+    with tf.variable_scope('branch_{}'.format(branch)):
+      b = _shake_shake_branch(x, output_filters, stride, r_forward, r_backward,
+                              is_training)
+      branches.append(b)
+  res = _shake_shake_skip_connection(x, output_filters, stride)
+  return res + tf.add_n(branches)
+
+
+def _shake_shake_layer(x, output_filters, num_blocks, stride,
+                       is_training):
+  """Builds many sub layers into one full layer."""
+  for block_num in range(num_blocks):
+    curr_stride = stride if (block_num == 0) else 1
+    with tf.variable_scope('layer_{}'.format(block_num)):
+      x = _shake_shake_block(x, output_filters, curr_stride,
+                             is_training)
+  return x
+
+
+def build_shake_shake_model(images, num_classes, hparams, is_training):
+  """Builds the Shake-Shake model.
+
+  Build the Shake-Shake model from https://arxiv.org/abs/1705.07485.
+
+  Args:
+    images: Tensor of images that will be fed into the Wide ResNet Model.
+    num_classes: Number of classed that the model needs to predict.
+    hparams: tf.HParams object that contains additional hparams needed to
+      construct the model. In this case it is the `shake_shake_widen_factor`
+      that is used to determine how many filters the model has.
+    is_training: Is the model training or not.
+
+  Returns:
+    The logits of the Shake-Shake model.
+  """
+  depth = 26
+  k = hparams.shake_shake_widen_factor  # The widen factor
+  n = int((depth - 2) / 6)
+  x = images
+
+  x = ops.conv2d(x, 16, 3, scope='init_conv')
+  x = ops.batch_norm(x, scope='init_bn')
+  with tf.variable_scope('L1'):
+    x = _shake_shake_layer(x, 16 * k, n, 1, is_training)
+  with tf.variable_scope('L2'):
+    x = _shake_shake_layer(x, 32 * k, n, 2, is_training)
+  with tf.variable_scope('L3'):
+    x = _shake_shake_layer(x, 64 * k, n, 2, is_training)
+  x = tf.nn.relu(x)
+  x = ops.global_avg_pool(x)
+
+  # Fully connected
+  logits = ops.fc(x, num_classes)
+  return logits

research/autoaugment/train_cifar.py

+# Copyright 2018 The TensorFlow Authors All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""AutoAugment Train/Eval module.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import contextlib
+import os
+import time
+
+import custom_ops as ops
+import data_utils
+import helper_utils
+import numpy as np
+from shake_drop import build_shake_drop_model
+from shake_shake import build_shake_shake_model
+import tensorflow as tf
+from wrn import build_wrn_model
+
+tf.flags.DEFINE_string('model_name', 'wrn',
+                       'wrn, shake_shake_32, shake_shake_96, shake_shake_112, '
+                       'pyramid_net')
+tf.flags.DEFINE_string('checkpoint_dir', '/tmp/training', 'Training Directory.')
+tf.flags.DEFINE_string('data_path', '/tmp/data',
+                       'Directory where dataset is located.')
+tf.flags.DEFINE_string('dataset', 'cifar10',
+                       'Dataset to train with. Either cifar10 or cifar100')
+tf.flags.DEFINE_integer('use_cpu', 1, '1 if use CPU, else GPU.')
+
+FLAGS = tf.flags.FLAGS
+
+arg_scope = tf.contrib.framework.arg_scope
+
+
+def setup_arg_scopes(is_training):
+  """Sets up the argscopes that will be used when building an image model.
+
+  Args:
+    is_training: Is the model training or not.
+
+  Returns:
+    Arg scopes to be put around the model being constructed.
+  """
+
+  batch_norm_decay = 0.9
+  batch_norm_epsilon = 1e-5
+  batch_norm_params = {
+      # Decay for the moving averages.
+      'decay': batch_norm_decay,
+      # epsilon to prevent 0s in variance.
+      'epsilon': batch_norm_epsilon,
+      'scale': True,
+      # collection containing the moving mean and moving variance.
+      'is_training': is_training,
+  }
+
+  scopes = []
+
+  scopes.append(arg_scope([ops.batch_norm], **batch_norm_params))
+  return scopes
+
+
+def build_model(inputs, num_classes, is_training, hparams):
+  """Constructs the vision model being trained/evaled.
+
+  Args:
+    inputs: input features/images being fed to the image model build built.
+    num_classes: number of output classes being predicted.
+    is_training: is the model training or not.
+    hparams: additional hyperparameters associated with the image model.
+
+  Returns:
+    The logits of the image model.
+  """
+  scopes = setup_arg_scopes(is_training)
+  with contextlib.nested(*scopes):
+    if hparams.model_name == 'pyramid_net':
+      logits = build_shake_drop_model(
+          inputs, num_classes, is_training)
+    elif hparams.model_name == 'wrn':
+      logits = build_wrn_model(
+          inputs, num_classes, hparams.wrn_size)
+    elif hparams.model_name == 'shake_shake':
+      logits = build_shake_shake_model(
+          inputs, num_classes, hparams, is_training)
+  return logits
+
+
+class CifarModel(object):
+  """Builds an image model for Cifar10/Cifar100."""
+
+  def __init__(self, hparams):
+    self.hparams = hparams
+
+  def build(self, mode):
+    """Construct the cifar model."""
+    assert mode in ['train', 'eval']
+    self.mode = mode
+    self._setup_misc(mode)
+    self._setup_images_and_labels()
+    self._build_graph(self.images, self.labels, mode)
+
+    self.init = tf.group(tf.global_variables_initializer(),
+                         tf.local_variables_initializer())
+
+  def _setup_misc(self, mode):
+    """Sets up miscellaneous in the cifar model constructor."""
+    self.lr_rate_ph = tf.Variable(0.0, name='lrn_rate', trainable=False)
+    self.reuse = None if (mode == 'train') else True
+    self.batch_size = self.hparams.batch_size
+    if mode == 'eval':
+      self.batch_size = 25
+
+  def _setup_images_and_labels(self):
+    """Sets up image and label placeholders for the cifar model."""
+    if FLAGS.dataset == 'cifar10':
+      self.num_classes = 10
+    else:
+      self.num_classes = 100
+    self.images = tf.placeholder(tf.float32, [self.batch_size, 32, 32, 3])
+    self.labels = tf.placeholder(tf.float32,
+                                 [self.batch_size, self.num_classes])
+
+  def assign_epoch(self, session, epoch_value):
+    session.run(self._epoch_update, feed_dict={self._new_epoch: epoch_value})
+
+  def _build_graph(self, images, labels, mode):
+    """Constructs the TF graph for the cifar model.
+
+    Args:
+      images: A 4-D image Tensor
+      labels: A 2-D labels Tensor.
+      mode: string indicating training mode ( e.g., 'train', 'valid', 'test').
+    """
+    is_training = 'train' in mode
+    if is_training:
+      self.global_step = tf.train.get_or_create_global_step()
+
+    logits = build_model(
+        images,
+        self.num_classes,
+        is_training,
+        self.hparams)
+    self.predictions, self.cost = helper_utils.setup_loss(
+        logits, labels)
+    self.accuracy, self.eval_op = tf.metrics.accuracy(
+        tf.argmax(labels, 1), tf.argmax(self.predictions, 1))
+    self._calc_num_trainable_params()
+
+    # Adds L2 weight decay to the cost
+    self.cost = helper_utils.decay_weights(self.cost,
+                                           self.hparams.weight_decay_rate)
+
+    if is_training:
+      self._build_train_op()
+
+    # Setup checkpointing for this child model
+    # Keep 2 or more checkpoints around during training.
+    with tf.device('/cpu:0'):
+      self.saver = tf.train.Saver(max_to_keep=2)
+
+    self.init = tf.group(tf.global_variables_initializer(),
+                         tf.local_variables_initializer())
+
+  def _calc_num_trainable_params(self):
+    self.num_trainable_params = np.sum([
+        np.prod(var.get_shape().as_list()) for var in tf.trainable_variables()
+    ])
+    tf.logging.info('number of trainable params: {}'.format(
+        self.num_trainable_params))
+
+  def _build_train_op(self):
+    """Builds the train op for the cifar model."""
+    hparams = self.hparams
+    tvars = tf.trainable_variables()
+    grads = tf.gradients(self.cost, tvars)
+    if hparams.gradient_clipping_by_global_norm > 0.0:
+      grads, norm = tf.clip_by_global_norm(
+          grads, hparams.gradient_clipping_by_global_norm)
+      tf.summary.scalar('grad_norm', norm)
+
+    # Setup the initial learning rate
+    initial_lr = self.lr_rate_ph
+    optimizer = tf.train.MomentumOptimizer(
+        initial_lr,
+        0.9,
+        use_nesterov=True)
+
+    self.optimizer = optimizer
+    apply_op = optimizer.apply_gradients(
+        zip(grads, tvars), global_step=self.global_step, name='train_step')
+    train_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
+    with tf.control_dependencies([apply_op]):
+      self.train_op = tf.group(*train_ops)
+
+
+class CifarModelTrainer(object):
+  """Trains an instance of the CifarModel class."""
+
+  def __init__(self, hparams):
+    self._session = None
+    self.hparams = hparams
+
+    self.model_dir = os.path.join(FLAGS.checkpoint_dir, 'model')
+    self.log_dir = os.path.join(FLAGS.checkpoint_dir, 'log')
+    # Set the random seed to be sure the same validation set
+    # is used for each model
+    np.random.seed(0)
+    self.data_loader = data_utils.DataSet(hparams)
+    np.random.seed()  # Put the random seed back to random
+    self.data_loader.reset()
+
+  def save_model(self, step=None):
+    """Dumps model into the backup_dir.
+
+    Args:
+      step: If provided, creates a checkpoint with the given step
+        number, instead of overwriting the existing checkpoints.
+    """
+    model_save_name = os.path.join(self.model_dir, 'model.ckpt')
+    if not tf.gfile.IsDirectory(self.model_dir):
+      tf.gfile.MakeDirs(self.model_dir)
+    self.saver.save(self.session, model_save_name, global_step=step)
+    tf.logging.info('Saved child model')
+
+  def extract_model_spec(self):
+    """Loads a checkpoint with the architecture structure stored in the name."""
+    checkpoint_path = tf.train.latest_checkpoint(self.model_dir)
+    if checkpoint_path is not None:
+      self.saver.restore(self.session, checkpoint_path)
+      tf.logging.info('Loaded child model checkpoint from %s',
+                      checkpoint_path)
+    else:
+      self.save_model(step=0)
+
+  def eval_child_model(self, model, data_loader, mode):
+    """Evaluate the child model.
+
+    Args:
+      model: image model that will be evaluated.
+      data_loader: dataset object to extract eval data from.
+      mode: will the model be evalled on train, val or test.
+
+    Returns:
+      Accuracy of the model on the specified dataset.
+    """
+    tf.logging.info('Evaluating child model in mode %s', mode)
+    while True:
+      try:
+        with self._new_session(model):
+          accuracy = helper_utils.eval_child_model(
+              self.session,
+              model,
+              data_loader,
+              mode)
+          tf.logging.info('Eval child model accuracy: {}'.format(accuracy))
+          # If epoch trained without raising the below errors, break
+          # from loop.
+          break
+      except (tf.errors.AbortedError, tf.errors.UnavailableError) as e:
+        tf.logging.info('Retryable error caught: %s.  Retrying.', e)
+
+    return accuracy
+
+  @contextlib.contextmanager
+  def _new_session(self, m):
+    """Creates a new session for model m."""
+    # Create a new session for this model, initialize
+    # variables, and save / restore from
+    # checkpoint.
+    self._session = tf.Session(
+        '',
+        config=tf.ConfigProto(
+            allow_soft_placement=True, log_device_placement=False))
+    self.session.run(m.init)
+
+    # Load in a previous checkpoint, or save this one
+    self.extract_model_spec()
+    try:
+      yield
+    finally:
+      tf.Session.reset('')
+      self._session = None
+
+  def _build_models(self):
+    """Builds the image models for train and eval."""
+    # Determine if we should build the train and eval model. When using
+    # distributed training we only want to build one or the other and not both.
+    with tf.variable_scope('model', use_resource=False):
+      m = CifarModel(self.hparams)
+      m.build('train')
+      self._num_trainable_params = m.num_trainable_params
+      self._saver = m.saver
+    with tf.variable_scope('model', reuse=True, use_resource=False):
+      meval = CifarModel(self.hparams)
+      meval.build('eval')
+    return m, meval
+
+  def _calc_starting_epoch(self, m):
+    """Calculates the starting epoch for model m based on global step."""
+    hparams = self.hparams
+    batch_size = hparams.batch_size
+    steps_per_epoch = int(hparams.train_size / batch_size)
+    with self._new_session(m):
+      curr_step = self.session.run(m.global_step)
+    total_steps = steps_per_epoch * hparams.num_epochs
+    epochs_left = (total_steps - curr_step) // steps_per_epoch
+    starting_epoch = hparams.num_epochs - epochs_left
+    return starting_epoch
+
+  def _run_training_loop(self, m, curr_epoch):
+    """Trains the cifar model `m` for one epoch."""
+    start_time = time.time()
+    while True:
+      try:
+        with self._new_session(m):
+          train_accuracy = helper_utils.run_epoch_training(
+              self.session, m, self.data_loader, curr_epoch)
+          tf.logging.info('Saving model after epoch')
+          self.save_model(step=curr_epoch)
+          break
+      except (tf.errors.AbortedError, tf.errors.UnavailableError) as e:
+        tf.logging.info('Retryable error caught: %s.  Retrying.', e)
+    tf.logging.info('Finished epoch: {}'.format(curr_epoch))
+    tf.logging.info('Epoch time(min): {}'.format(
+        (time.time() - start_time) / 60.0))
+    return train_accuracy
+
+  def _compute_final_accuracies(self, meval):
+    """Run once training is finished to compute final val/test accuracies."""
+    valid_accuracy = self.eval_child_model(meval, self.data_loader, 'val')
+    if self.hparams.eval_test:
+      test_accuracy = self.eval_child_model(meval, self.data_loader, 'test')
+    else:
+      test_accuracy = 0
+    tf.logging.info('Test Accuracy: {}'.format(test_accuracy))
+    return valid_accuracy, test_accuracy
+
+  def run_model(self):
+    """Trains and evalutes the image model."""
+    hparams = self.hparams
+
+    # Build the child graph
+    with tf.Graph().as_default(), tf.device(
+        '/cpu:0' if FLAGS.use_cpu else '/gpu:0'):
+      m, meval = self._build_models()
+
+      # Figure out what epoch we are on
+      starting_epoch = self._calc_starting_epoch(m)
+
+      # Run the validation error right at the beginning
+      valid_accuracy = self.eval_child_model(
+          meval, self.data_loader, 'val')
+      tf.logging.info('Before Training Epoch: {}     Val Acc: {}'.format(
+          starting_epoch, valid_accuracy))
+      training_accuracy = None
+
+      for curr_epoch in xrange(starting_epoch, hparams.num_epochs):
+
+        # Run one training epoch
+        training_accuracy = self._run_training_loop(m, curr_epoch)
+
+        valid_accuracy = self.eval_child_model(
+            meval, self.data_loader, 'val')
+        tf.logging.info('Epoch: {}    Valid Acc: {}'.format(
+            curr_epoch, valid_accuracy))
+
+      valid_accuracy, test_accuracy = self._compute_final_accuracies(
+          meval)
+
+    tf.logging.info(
+        'Train Acc: {}    Valid Acc: {}     Test Acc: {}'.format(
+            training_accuracy, valid_accuracy, test_accuracy))
+
+  @property
+  def saver(self):
+    return self._saver
+
+  @property
+  def session(self):
+    return self._session
+
+  @property
+  def num_trainable_params(self):
+    return self._num_trainable_params
+
+
+def main(_):
+  if FLAGS.dataset not in ['cifar10', 'cifar100']:
+    raise ValueError('Invalid dataset: %s' % FLAGS.dataset)
+  hparams = tf.contrib.training.HParams(
+      train_size=50000,
+      validation_size=0,
+      eval_test=1,
+      dataset=FLAGS.dataset,
+      data_path=FLAGS.data_path,
+      batch_size=128,
+      gradient_clipping_by_global_norm=5.0)
+  if FLAGS.model_name == 'wrn':
+    hparams.add_hparam('model_name', 'wrn')
+    hparams.add_hparam('num_epochs', 200)
+    hparams.add_hparam('wrn_size', 160)
+    hparams.add_hparam('lr', 0.1)
+    hparams.add_hparam('weight_decay_rate', 5e-4)
+  elif FLAGS.model_name == 'shake_shake_32':
+    hparams.add_hparam('model_name', 'shake_shake')
+    hparams.add_hparam('num_epochs', 1800)
+    hparams.add_hparam('shake_shake_widen_factor', 2)
+    hparams.add_hparam('lr', 0.01)
+    hparams.add_hparam('weight_decay_rate', 0.001)
+  elif FLAGS.model_name == 'shake_shake_96':
+    hparams.add_hparam('model_name', 'shake_shake')
+    hparams.add_hparam('num_epochs', 1800)
+    hparams.add_hparam('shake_shake_widen_factor', 6)
+    hparams.add_hparam('lr', 0.01)
+    hparams.add_hparam('weight_decay_rate', 0.001)
+  elif FLAGS.model_name == 'shake_shake_112':
+    hparams.add_hparam('model_name', 'shake_shake')
+    hparams.add_hparam('num_epochs', 1800)
+    hparams.add_hparam('shake_shake_widen_factor', 7)
+    hparams.add_hparam('lr', 0.01)
+    hparams.add_hparam('weight_decay_rate', 0.001)
+  elif FLAGS.model_name == 'pyramid_net':
+    hparams.add_hparam('model_name', 'pyramid_net')
+    hparams.add_hparam('num_epochs', 1800)
+    hparams.add_hparam('lr', 0.05)
+    hparams.add_hparam('weight_decay_rate', 5e-5)
+    hparams.batch_size = 64
+  else:
+    raise ValueError('Not Valid Model Name: %s' % FLAGS.model_name)
+  cifar_trainer = CifarModelTrainer(hparams)
+  cifar_trainer.run_model()
+
+if __name__ == '__main__':
+  tf.logging.set_verbosity(tf.logging.INFO)
+  tf.app.run()

research/autoaugment/wrn.py

+# Copyright 2018 The TensorFlow Authors All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Builds the Wide-ResNet Model."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import custom_ops as ops
+import numpy as np
+import tensorflow as tf
+
+
+
+def residual_block(
+    x, in_filter, out_filter, stride, activate_before_residual=False):
+  """Adds residual connection to `x` in addition to applying BN->ReLU->3x3 Conv.
+
+  Args:
+    x: Tensor that is the output of the previous layer in the model.
+    in_filter: Number of filters `x` has.
+    out_filter: Number of filters that the output of this layer will have.
+    stride: Integer that specified what stride should be applied to `x`.
+    activate_before_residual: Boolean on whether a BN->ReLU should be applied
+      to x before the convolution is applied.
+
+  Returns:
+    A Tensor that is the result of applying two sequences of BN->ReLU->3x3 Conv
+    and then adding that Tensor to `x`.
+  """
+
+  if activate_before_residual:  # Pass up RELU and BN activation for resnet
+    with tf.variable_scope('shared_activation'):
+      x = ops.batch_norm(x, scope='init_bn')
+      x = tf.nn.relu(x)
+      orig_x = x
+  else:
+    orig_x = x
+
+  block_x = x
+  if not activate_before_residual:
+    with tf.variable_scope('residual_only_activation'):
+      block_x = ops.batch_norm(block_x, scope='init_bn')
+      block_x = tf.nn.relu(block_x)
+
+  with tf.variable_scope('sub1'):
+    block_x = ops.conv2d(
+        block_x, out_filter, 3, stride=stride, scope='conv1')
+
+  with tf.variable_scope('sub2'):
+    block_x = ops.batch_norm(block_x, scope='bn2')
+    block_x = tf.nn.relu(block_x)
+    block_x = ops.conv2d(
+        block_x, out_filter, 3, stride=1, scope='conv2')
+
+  with tf.variable_scope(
+      'sub_add'):  # If number of filters do not agree then zero pad them
+    if in_filter != out_filter:
+      orig_x = ops.avg_pool(orig_x, stride, stride)
+      orig_x = ops.zero_pad(orig_x, in_filter, out_filter)
+  x = orig_x + block_x
+  return x
+
+
+def _res_add(in_filter, out_filter, stride, x, orig_x):
+  """Adds `x` with `orig_x`, both of which are layers in the model.
+
+  Args:
+    in_filter: Number of filters in `orig_x`.
+    out_filter: Number of filters in `x`.
+    stride: Integer specifying the stide that should be applied `orig_x`.
+    x: Tensor that is the output of the previous layer.
+    orig_x: Tensor that is the output of an earlier layer in the network.
+
+  Returns:
+    A Tensor that is the result of `x` and `orig_x` being added after
+    zero padding and striding are applied to `orig_x` to get the shapes
+    to match.
+  """
+  if in_filter != out_filter:
+    orig_x = ops.avg_pool(orig_x, stride, stride)
+    orig_x = ops.zero_pad(orig_x, in_filter, out_filter)
+  x = x + orig_x
+  orig_x = x
+  return x, orig_x
+
+
+def build_wrn_model(images, num_classes, wrn_size):
+  """Builds the WRN model.
+
+  Build the Wide ResNet model from https://arxiv.org/abs/1605.07146.
+
+  Args:
+    images: Tensor of images that will be fed into the Wide ResNet Model.
+    num_classes: Number of classed that the model needs to predict.
+    wrn_size: Parameter that scales the number of filters in the Wide ResNet
+      model.
+
+  Returns:
+    The logits of the Wide ResNet model.
+  """
+  kernel_size = wrn_size
+  filter_size = 3
+  num_blocks_per_resnet = 4
+  filters = [
+      min(kernel_size, 16), kernel_size, kernel_size * 2, kernel_size * 4
+  ]
+  strides = [1, 2, 2]  # stride for each resblock
+
+  # Run the first conv
+  with tf.variable_scope('init'):
+    x = images
+    output_filters = filters[0]
+    x = ops.conv2d(x, output_filters, filter_size, scope='init_conv')
+
+  first_x = x  # Res from the beginning
+  orig_x = x  # Res from previous block
+
+  for block_num in range(1, 4):
+    with tf.variable_scope('unit_{}_0'.format(block_num)):
+      activate_before_residual = True if block_num == 1 else False
+      x = residual_block(
+          x,
+          filters[block_num - 1],
+          filters[block_num],
+          strides[block_num - 1],
+          activate_before_residual=activate_before_residual)
+    for i in range(1, num_blocks_per_resnet):
+      with tf.variable_scope('unit_{}_{}'.format(block_num, i)):
+        x = residual_block(
+            x,
+            filters[block_num],
+            filters[block_num],
+            1,
+            activate_before_residual=False)
+    x, orig_x = _res_add(filters[block_num - 1], filters[block_num],
+                         strides[block_num - 1], x, orig_x)
+  final_stride_val = np.prod(strides)
+  x, _ = _res_add(filters[0], filters[3], final_stride_val, x, first_x)
+  with tf.variable_scope('unit_last'):
+    x = ops.batch_norm(x, scope='final_bn')
+    x = tf.nn.relu(x)
+    x = ops.global_avg_pool(x)
+    logits = ops.fc(x, num_classes)
+  return logits

research/cognitive_planning/BUILD

+package(default_visibility = [":internal"])
+
+licenses(["notice"])  # Apache 2.0
+
+exports_files(["LICENSE"])
+
+package_group(
+    name = "internal",
+    packages = [
+        "//cognitive_planning/...",
+    ],
+)
+
+py_binary(
+    name = "train_supervised_active_vision",
+    srcs = [
+        "train_supervised_active_vision.py",
+    ],
+)

research/cognitive_planning/README.md

+# cognitive_planning
+
+**Visual Representation for Semantic Target Driven Navigation**
+
+Arsalan Mousavian, Alexander Toshev, Marek Fiser, Jana Kosecka, James Davidson
+
+This is the implementation of semantic target driven navigation training and evaluation on 
+Active Vision dataset. 
+
+ECCV Workshop on Visual Learning and Embodied Agents in Simulation Environments
+2018.
+
+<div align="center">
+  <table style="width:100%" border="0">
+    <tr>
+      <td align="center"><img src='https://cs.gmu.edu/~amousavi/gifs/smaller_fridge_2.gif'></td>
+      <td align="center"><img src='https://cs.gmu.edu/~amousavi/gifs/smaller_tv_1.gif'></td>
+    </tr>
+    <tr>
+      <td align="center">Target: Fridge</td>
+      <td align="center">Target: Television</td>
+    </tr>
+    <tr>
+      <td align="center"><img src='https://cs.gmu.edu/~amousavi/gifs/smaller_microwave_1.gif'></td>
+      <td align="center"><img src='https://cs.gmu.edu/~amousavi/gifs/smaller_couch_1.gif'></td>
+    </tr>
+    <tr>
+      <td align="center">Target: Microwave</td>
+      <td align="center">Target: Couch</td>
+    </tr>
+  </table>
+</div>
+
+
+
+Paper: [https://arxiv.org/abs/1805.06066](https://arxiv.org/abs/1805.06066)
+
+
+## 1. Installation
+
+### Requirements
+
+#### Python Packages
+
+```shell
+networkx
+gin-config
+```
+
+### Download cognitive_planning
+
+```shell
+git clone --depth 1 https://github.com/tensorflow/models.git
+```
+
+## 2. Datasets
+
+### Download ActiveVision Dataset 
+We used Active Vision Dataset (AVD) which can be downloaded from [here](http://cs.unc.edu/~ammirato/active_vision_dataset_website/). To make our code faster and reduce memory footprint, we created the AVD Minimal dataset. AVD Minimal consists of low resolution images from the original AVD dataset. In addition, we added annotations for target views, predicted object detections from pre-trained object detector on MS-COCO dataset, and predicted semantic segmentation from pre-trained model on NYU-v2 dataset. AVD minimal can be downloaded from [here](https://storage.googleapis.com/active-vision-dataset/AVD_Minimal.zip). Set `$AVD_DIR` as the path to the downloaded AVD Minimal.
+
+### TODO: SUNCG Dataset
+Current version of the code does not support SUNCG dataset. It can be added by
+implementing necessary functions of `envs/task_env.py` using the public
+released code of SUNCG environment such as
+[House3d](https://github.com/facebookresearch/House3D) and
+[MINOS](https://github.com/minosworld/minos). 
+
+### ActiveVisionDataset Demo
+
+
+If you wish to navigate the environment, to see how the AVD looks like you can use the following command:
+```shell
+python viz_active_vision_dataset_main -- \
+  --mode=human \
+  --gin_config=envs/configs/active_vision_config.gin \
+  --gin_params='ActiveVisionDatasetEnv.dataset_root=$AVD_DIR'
+```
+
+## 3. Training
+Right now, the released version only supports training and inference using the real data from Active Vision Dataset.
+
+When RGB image modality is used, the Resnet embeddings are initialized. To start the training download pre-trained Resnet50 check point in the working directory ./resnet_v2_50_checkpoint/resnet_v2_50.ckpt
+
+```
+wget http://download.tensorflow.org/models/resnet_v2_50_2017_04_14.tar.gz
+```
+### Run training
+Use the following command for training:
+```shell
+# Train
+python train_supervised_active_vision.py \
+  --mode='train' \
+  --logdir=$CHECKPOINT_DIR \
+  --modality_types='det' \
+  --batch_size=8 \
+  --train_iters=200000 \
+  --lstm_cell_size=2048 \
+  --policy_fc_size=2048 \
+  --sequence_length=20 \
+  --max_eval_episode_length=100 \
+  --test_iters=194 \
+  --gin_config=envs/configs/active_vision_config.gin \
+  --gin_params='ActiveVisionDatasetEnv.dataset_root=$AVD_DIR' \
+  --logtostderr
+```
+
+The training can be run for different modalities and modality combinations, including semantic segmentation, object detectors, RGB images, depth images. Low resolution images and outputs of detectors pretrained on COCO dataset and semantic segmenation pre trained on NYU dataset are provided as a part of this distribution and can be found in Meta directory of AVD_Minimal. 
+Additional details are described in the comments of the code and in the paper.
+
+### Run Evaluation
+Use the following command for unrolling the policy on the eval environments. The inference code periodically check the checkpoint folder for new checkpoints to use it for unrolling the policy on the eval environments. After each evaluation, it will create a folder in the $CHECKPOINT_DIR/evals/$ITER where $ITER is the iteration number at which the checkpoint is stored.
+```shell
+# Eval
+python train_supervised_active_vision.py \
+  --mode='eval' \
+  --logdir=$CHECKPOINT_DIR \
+  --modality_types='det' \
+  --batch_size=8 \
+  --train_iters=200000 \
+  --lstm_cell_size=2048 \
+  --policy_fc_size=2048 \
+  --sequence_length=20 \
+  --max_eval_episode_length=100 \
+  --test_iters=194 \
+  --gin_config=envs/configs/active_vision_config.gin \
+  --gin_params='ActiveVisionDatasetEnv.dataset_root=$AVD_DIR' \
+  --logtostderr
+```
+At any point, you can run the following command to compute statistics such as success rate over all the evaluations so far. It also generates gif images for unrolling of the best policy.
+```shell
+# Visualize and Compute Stats
+python viz_active_vision_dataset_main.py \
+   --mode=eval \ 
+   --eval_folder=$CHECKPOINT_DIR/evals/ \
+   --output_folder=$OUTPUT_GIFS_FOLDER \
+   --gin_config=envs/configs/active_vision_config.gin \
+   --gin_params='ActiveVisionDatasetEnv.dataset_root=$AVD_DIR'
+```
+## Contact
+
+To ask questions or report issues please open an issue on the tensorflow/models
+[issues tracker](https://github.com/tensorflow/models/issues).
+Please assign issues to @arsalan-mousavian.
+
+## Reference
+The details of the training and experiments can be found in the following paper. If you find our work useful in your research please consider citing our paper:
+
+```
+@inproceedings{MousavianECCVW18,
+  author = {A. Mousavian and A. Toshev and M. Fiser and J. Kosecka and J. Davidson},
+  title = {Visual Representations for Semantic Target Driven Navigation},
+  booktitle = {ECCV Workshop on Visual Learning and Embodied Agents in Simulation Environments},
+  year = {2018},
+}
+```
+
+

research/cognitive_planning/command

+python train_supervised_active_vision \
+  --mode='train' \
+  --logdir=/usr/local/google/home/kosecka/checkin_log_det/ \
+  --modality_types='det' \
+  --batch_size=8 \
+  --train_iters=200000 \
+  --lstm_cell_size=2048 \
+  --policy_fc_size=2048 \
+  --sequence_length=20 \
+  --max_eval_episode_length=100 \
+  --test_iters=194 \
+  --gin_config=robotics/cognitive_planning/envs/configs/active_vision_config.gin \
+  --gin_params='ActiveVisionDatasetEnv.dataset_root="/usr/local/google/home/kosecka/AVD_minimal/"' \
+  --logtostderr

research/cognitive_planning/embedders.py

+# Copyright 2018 The TensorFlow Authors All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Interface for different embedders for modalities."""
+
+import abc
+import numpy as np
+import tensorflow as tf
+import preprocessing
+from tensorflow.contrib.slim.nets import resnet_v2
+
+slim = tf.contrib.slim
+
+
+class Embedder(object):
+  """Represents the embedder for different modalities.
+
+  Modalities can be semantic segmentation, depth channel, object detection and
+  so on, which require specific embedder for them.
+  """
+  __metaclass__ = abc.ABCMeta
+
+  @abc.abstractmethod
+  def build(self, observation):
+    """Builds the model to embed the observation modality.
+
+    Args:
+      observation: tensor that contains the raw observation from modality.
+    Returns:
+      Embedding tensor for the given observation tensor.
+    """
+    raise NotImplementedError(
+        'Needs to be implemented as part of Embedder Interface')
+
+
+class DetectionBoxEmbedder(Embedder):
+  """Represents the model that encodes the detection boxes from images."""
+
+  def __init__(self, rnn_state_size, scope=None):
+    self._rnn_state_size = rnn_state_size
+    self._scope = scope
+
+  def build(self, observations):
+    """Builds the model to embed object detection observations.
+
+    Args:
+      observations: a tuple of (dets, det_num).
+        dets is a tensor of BxTxLxE that has the detection boxes in all the
+          images of the batch. B is the batch size, T is the maximum length of
+          episode, L is the maximum number of detections per image in the batch
+          and E is the size of each detection embedding.
+        det_num is a tensor of BxT that contains the number of detected boxes
+          each image of each sequence in the batch.
+    Returns:
+      For each image in the batch, returns the accumulative embedding of all the
+      detection boxes in that image.
+    """
+    with tf.variable_scope(self._scope, default_name=''):
+      shape = observations[0].shape
+      dets = tf.reshape(observations[0], [-1, shape[-2], shape[-1]])
+      det_num = tf.reshape(observations[1], [-1])
+      lstm_cell = tf.nn.rnn_cell.BasicLSTMCell(self._rnn_state_size)
+      batch_size = tf.shape(dets)[0]
+      lstm_outputs, _ = tf.nn.dynamic_rnn(
+          cell=lstm_cell,
+          inputs=dets,
+          sequence_length=det_num,
+          initial_state=lstm_cell.zero_state(batch_size, dtype=tf.float32),
+          dtype=tf.float32)
+      # Gathering the last state of each sequence in the batch.
+      batch_range = tf.range(batch_size)
+      indices = tf.stack([batch_range, det_num - 1], axis=1)
+      last_lstm_outputs = tf.gather_nd(lstm_outputs, indices)
+      last_lstm_outputs = tf.reshape(last_lstm_outputs,
+                                     [-1, shape[1], self._rnn_state_size])
+    return last_lstm_outputs
+
+
+class ResNet(Embedder):
+  """Residual net embedder for image data."""
+
+  def __init__(self, params, *args, **kwargs):
+    super(ResNet, self).__init__(*args, **kwargs)
+    self._params = params
+    self._extra_train_ops = []
+
+  def build(self, images):
+    shape = images.get_shape().as_list()
+    if len(shape) == 5:
+      images = tf.reshape(images,
+                          [shape[0] * shape[1], shape[2], shape[3], shape[4]])
+    embedding = self._build_model(images)
+    if len(shape) == 5:
+      embedding = tf.reshape(embedding, [shape[0], shape[1], -1])
+
+    return embedding
+
+  @property
+  def extra_train_ops(self):
+    return self._extra_train_ops
+
+  def _build_model(self, images):
+    """Builds the model."""
+
+    # Convert images to floats and normalize them.
+    images = tf.to_float(images)
+    bs = images.get_shape().as_list()[0]
+    images = [
+        tf.image.per_image_standardization(tf.squeeze(i))
+        for i in tf.split(images, bs)
+    ]
+    images = tf.concat([tf.expand_dims(i, axis=0) for i in images], axis=0)
+
+    with tf.variable_scope('init'):
+      x = self._conv('init_conv', images, 3, 3, 16, self._stride_arr(1))
+
+    strides = [1, 2, 2]
+    activate_before_residual = [True, False, False]
+    if self._params.use_bottleneck:
+      res_func = self._bottleneck_residual
+      filters = [16, 64, 128, 256]
+    else:
+      res_func = self._residual
+      filters = [16, 16, 32, 128]
+
+    with tf.variable_scope('unit_1_0'):
+      x = res_func(x, filters[0], filters[1], self._stride_arr(strides[0]),
+                   activate_before_residual[0])
+    for i in xrange(1, self._params.num_residual_units):
+      with tf.variable_scope('unit_1_%d' % i):
+        x = res_func(x, filters[1], filters[1], self._stride_arr(1), False)
+
+    with tf.variable_scope('unit_2_0'):
+      x = res_func(x, filters[1], filters[2], self._stride_arr(strides[1]),
+                   activate_before_residual[1])
+    for i in xrange(1, self._params.num_residual_units):
+      with tf.variable_scope('unit_2_%d' % i):
+        x = res_func(x, filters[2], filters[2], self._stride_arr(1), False)
+
+    with tf.variable_scope('unit_3_0'):
+      x = res_func(x, filters[2], filters[3], self._stride_arr(strides[2]),
+                   activate_before_residual[2])
+    for i in xrange(1, self._params.num_residual_units):
+      with tf.variable_scope('unit_3_%d' % i):
+        x = res_func(x, filters[3], filters[3], self._stride_arr(1), False)
+
+    with tf.variable_scope('unit_last'):
+      x = self._batch_norm('final_bn', x)
+      x = self._relu(x, self._params.relu_leakiness)
+
+    with tf.variable_scope('pool_logit'):
+      x = self._global_avg_pooling(x)
+
+    return x
+
+  def _stride_arr(self, stride):
+    return [1, stride, stride, 1]
+
+  def _batch_norm(self, name, x):
+    """batch norm implementation."""
+    with tf.variable_scope(name):
+      params_shape = [x.shape[-1]]
+
+      beta = tf.get_variable(
+          'beta',
+          params_shape,
+          tf.float32,
+          initializer=tf.constant_initializer(0.0, tf.float32))
+      gamma = tf.get_variable(
+          'gamma',
+          params_shape,
+          tf.float32,
+          initializer=tf.constant_initializer(1.0, tf.float32))
+
+      if self._params.is_train:
+        mean, variance = tf.nn.moments(x, [0, 1, 2], name='moments')
+
+        moving_mean = tf.get_variable(
+            'moving_mean',
+            params_shape,
+            tf.float32,
+            initializer=tf.constant_initializer(0.0, tf.float32),
+            trainable=False)
+        moving_variance = tf.get_variable(
+            'moving_variance',
+            params_shape,
+            tf.float32,
+            initializer=tf.constant_initializer(1.0, tf.float32),
+            trainable=False)
+
+        self._extra_train_ops.append(
+            tf.assign_moving_average(moving_mean, mean, 0.9))
+        self._extra_train_ops.append(
+            tf.assign_moving_average(moving_variance, variance, 0.9))
+      else:
+        mean = tf.get_variable(
+            'moving_mean',
+            params_shape,
+            tf.float32,
+            initializer=tf.constant_initializer(0.0, tf.float32),
+            trainable=False)
+        variance = tf.get_variable(
+            'moving_variance',
+            params_shape,
+            tf.float32,
+            initializer=tf.constant_initializer(1.0, tf.float32),
+            trainable=False)
+        tf.summary.histogram(mean.op.name, mean)
+        tf.summary.histogram(variance.op.name, variance)
+      # elipson used to be 1e-5. Maybe 0.001 solves NaN problem in deeper net.
+      y = tf.nn.batch_normalization(x, mean, variance, beta, gamma, 0.001)
+      y.set_shape(x.shape)
+      return y
+
+  def _residual(self,
+                x,
+                in_filter,
+                out_filter,
+                stride,
+                activate_before_residual=False):
+    """Residual unit with 2 sub layers."""
+
+    if activate_before_residual:
+      with tf.variable_scope('shared_activation'):
+        x = self._batch_norm('init_bn', x)
+        x = self._relu(x, self._params.relu_leakiness)
+        orig_x = x
+    else:
+      with tf.variable_scope('residual_only_activation'):
+        orig_x = x
+        x = self._batch_norm('init_bn', x)
+        x = self._relu(x, self._params.relu_leakiness)
+
+    with tf.variable_scope('sub1'):
+      x = self._conv('conv1', x, 3, in_filter, out_filter, stride)
+
+    with tf.variable_scope('sub2'):
+      x = self._batch_norm('bn2', x)
+      x = self._relu(x, self._params.relu_leakiness)
+      x = self._conv('conv2', x, 3, out_filter, out_filter, [1, 1, 1, 1])
+
+    with tf.variable_scope('sub_add'):
+      if in_filter != out_filter:
+        orig_x = tf.nn.avg_pool(orig_x, stride, stride, 'VALID')
+        orig_x = tf.pad(
+            orig_x, [[0, 0], [0, 0], [0, 0], [(out_filter - in_filter) // 2,
+                                              (out_filter - in_filter) // 2]])
+      x += orig_x
+
+    return x
+
+  def _bottleneck_residual(self,
+                           x,
+                           in_filter,
+                           out_filter,
+                           stride,
+                           activate_before_residual=False):
+    """A residual convolutional layer with a bottleneck.
+
+    The layer is a composite of three convolutional layers with a ReLU non-
+    linearity and batch normalization after each linear convolution. The depth
+    if the second and third layer is out_filter / 4 (hence it is a bottleneck).
+
+    Args:
+      x: a float 4 rank Tensor representing the input to the layer.
+      in_filter: a python integer representing depth of the input.
+      out_filter: a python integer representing depth of the output.
+      stride: a python integer denoting the stride of the layer applied before
+        the first convolution.
+      activate_before_residual: a python boolean. If True, then a ReLU is
+        applied as a first operation on the input x before everything else.
+    Returns:
+      A 4 rank Tensor with batch_size = batch size of input, width and height =
+      width / stride and height / stride of the input and depth = out_filter.
+    """
+    if activate_before_residual:
+      with tf.variable_scope('common_bn_relu'):
+        x = self._batch_norm('init_bn', x)
+        x = self._relu(x, self._params.relu_leakiness)
+        orig_x = x
+    else:
+      with tf.variable_scope('residual_bn_relu'):
+        orig_x = x
+        x = self._batch_norm('init_bn', x)
+        x = self._relu(x, self._params.relu_leakiness)
+
+    with tf.variable_scope('sub1'):
+      x = self._conv('conv1', x, 1, in_filter, out_filter / 4, stride)
+
+    with tf.variable_scope('sub2'):
+      x = self._batch_norm('bn2', x)
+      x = self._relu(x, self._params.relu_leakiness)
+      x = self._conv('conv2', x, 3, out_filter / 4, out_filter / 4,
+                     [1, 1, 1, 1])
+
+    with tf.variable_scope('sub3'):
+      x = self._batch_norm('bn3', x)
+      x = self._relu(x, self._params.relu_leakiness)
+      x = self._conv('conv3', x, 1, out_filter / 4, out_filter, [1, 1, 1, 1])
+
+    with tf.variable_scope('sub_add'):
+      if in_filter != out_filter:
+        orig_x = self._conv('project', orig_x, 1, in_filter, out_filter, stride)
+      x += orig_x
+
+    return x
+
+  def _decay(self):
+    costs = []
+    for var in tf.trainable_variables():
+      if var.op.name.find(r'DW') > 0:
+        costs.append(tf.nn.l2_loss(var))
+
+    return tf.mul(self._params.weight_decay_rate, tf.add_n(costs))
+
+  def _conv(self, name, x, filter_size, in_filters, out_filters, strides):
+    """Convolution."""
+    with tf.variable_scope(name):
+      n = filter_size * filter_size * out_filters
+      kernel = tf.get_variable(
+          'DW', [filter_size, filter_size, in_filters, out_filters],
+          tf.float32,
+          initializer=tf.random_normal_initializer(stddev=np.sqrt(2.0 / n)))
+      return tf.nn.conv2d(x, kernel, strides, padding='SAME')
+
+  def _relu(self, x, leakiness=0.0):
+    return tf.where(tf.less(x, 0.0), leakiness * x, x, name='leaky_relu')
+
+  def _fully_connected(self, x, out_dim):
+    x = tf.reshape(x, [self._params.batch_size, -1])
+    w = tf.get_variable(
+        'DW', [x.get_shape()[1], out_dim],
+        initializer=tf.uniform_unit_scaling_initializer(factor=1.0))
+    b = tf.get_variable(
+        'biases', [out_dim], initializer=tf.constant_initializer())
+    return tf.nn.xw_plus_b(x, w, b)
+
+  def _global_avg_pooling(self, x):
+    assert x.get_shape().ndims == 4
+    return tf.reduce_mean(x, [1, 2])
+
+
+class MLPEmbedder(Embedder):
+  """Embedder of vectorial data.
+
+  The net is a multi-layer perceptron, with ReLU nonlinearities in all layers
+  except the last one.
+  """
+
+  def __init__(self, layers, *args, **kwargs):
+    """Constructs MLPEmbedder.
+
+    Args:
+      layers: a list of python integers representing layer sizes.
+      *args: arguments for super constructor.
+      **kwargs: keyed arguments for super constructor.
+    """
+    super(MLPEmbedder, self).__init__(*args, **kwargs)
+    self._layers = layers
+
+  def build(self, features):
+    shape = features.get_shape().as_list()
+    if len(shape) == 3:
+      features = tf.reshape(features, [shape[0] * shape[1], shape[2]])
+    x = features
+    for i, dim in enumerate(self._layers):
+      with tf.variable_scope('layer_%i' % i):
+        x = self._fully_connected(x, dim)
+        if i < len(self._layers) - 1:
+          x = self._relu(x)
+
+    if len(shape) == 3:
+      x = tf.reshape(x, shape[:-1] + [self._layers[-1]])
+    return x
+
+  def _fully_connected(self, x, out_dim):
+    w = tf.get_variable(
+        'DW', [x.get_shape()[1], out_dim],
+        initializer=tf.variance_scaling_initializer(distribution='uniform'))
+    b = tf.get_variable(
+        'biases', [out_dim], initializer=tf.constant_initializer())
+    return tf.nn.xw_plus_b(x, w, b)
+
+  def _relu(self, x, leakiness=0.0):
+    return tf.where(tf.less(x, 0.0), leakiness * x, x, name='leaky_relu')
+
+
+class SmallNetworkEmbedder(Embedder):
+  """Embedder for image like observations.
+
+  The network is comprised of multiple conv layers and a fully connected layer
+  at the end. The number of conv layers and the parameters are configured from
+  params.
+  """
+
+  def __init__(self, params, *args, **kwargs):
+    """Constructs the small network.
+
+    Args:
+      params: params should be tf.hparams type. params need to have a list of
+        conv_sizes, conv_strides, conv_channels. The length of these lists
+        should be equal to each other and to the number of conv layers in the
+        network. Plus, it also needs to have boolean variable named to_one_hot
+        which indicates whether the input should be converted to one hot or not.
+        The size of the fully connected layer is specified by
+        params.embedding_size.
+
+      *args: The rest of the parameters.
+      **kwargs: the reset of the parameters.
+
+    Raises:
+      ValueError: If the length of params.conv_strides, params.conv_sizes, and
+        params.conv_channels are not equal.
+
+    """
+
+    super(SmallNetworkEmbedder, self).__init__(*args, **kwargs)
+    self._params = params
+    if len(self._params.conv_sizes) != len(self._params.conv_strides):
+      raise ValueError(
+          'Conv sizes and strides should have the same length: {} != {}'.format(
+              len(self._params.conv_sizes), len(self._params.conv_strides)))
+
+    if len(self._params.conv_sizes) != len(self._params.conv_channels):
+      raise ValueError(
+          'Conv sizes and channels should have the same length: {} != {}'.
+          format(len(self._params.conv_sizes), len(self._params.conv_channels)))
+
+  def build(self, images):
+    """Builds the embedder with the given speicifcation.
+
+    Args:
+      images: a tensor that contains the input images which has the shape of
+        NxTxHxWxC where N is the batch size, T is the maximum length of the
+        sequence, H and W are the height and width of the images and C is the
+        number of channels.
+
+    Returns:
+      A tensor that is the embedding of the images.
+    """
+
+    shape = images.get_shape().as_list()
+    images = tf.reshape(images,
+                        [shape[0] * shape[1], shape[2], shape[3], shape[4]])
+
+    with slim.arg_scope(
+        [slim.conv2d, slim.fully_connected],
+        activation_fn=tf.nn.relu,
+        weights_regularizer=slim.l2_regularizer(self._params.weight_decay_rate),
+        biases_initializer=tf.zeros_initializer()):
+      with slim.arg_scope([slim.conv2d], padding='SAME'):
+        # convert the image to one hot if needed.
+        if self._params.to_one_hot:
+          net = tf.one_hot(
+              tf.squeeze(tf.to_int32(images), axis=[-1]),
+              self._params.one_hot_length)
+        else:
+          net = images
+
+        p = self._params
+        # Adding conv layers with the specified configurations.
+        for conv_id, kernel_stride_channel in enumerate(
+            zip(p.conv_sizes, p.conv_strides, p.conv_channels)):
+          kernel_size, stride, channels = kernel_stride_channel
+          net = slim.conv2d(
+              net,
+              channels, [kernel_size, kernel_size],
+              stride,
+              scope='conv_{}'.format(conv_id + 1))
+
+        net = slim.flatten(net)
+        net = slim.fully_connected(net, self._params.embedding_size, scope='fc')
+
+        output = tf.reshape(net, [shape[0], shape[1], -1])
+        return output
+
+
+class ResNet50Embedder(Embedder):
+  """Uses ResNet50 to embed input images."""
+
+  def build(self, images):
+    """Builds a ResNet50 embedder for the input images.
+
+    It assumes that the range of the pixel values in the images tensor is
+      [0,255] and should be castable to tf.uint8.
+
+    Args:
+      images: a tensor that contains the input images which has the shape of
+          NxTxHxWx3 where N is the batch size, T is the maximum length of the
+          sequence, H and W are the height and width of the images and C is the
+          number of channels.
+    Returns:
+      The embedding of the input image with the shape of NxTxL where L is the
+        embedding size of the output.
+
+    Raises:
+      ValueError: if the shape of the input does not agree with the expected
+      shape explained in the Args section.
+    """
+    shape = images.get_shape().as_list()
+    if len(shape) != 5:
+      raise ValueError(
+          'The tensor shape should have 5 elements, {} is provided'.format(
+              len(shape)))
+    if shape[4] != 3:
+      raise ValueError('Three channels are expected for the input image')
+
+    images = tf.cast(images, tf.uint8)
+    images = tf.reshape(images,
+                        [shape[0] * shape[1], shape[2], shape[3], shape[4]])
+    with slim.arg_scope(resnet_v2.resnet_arg_scope()):
+
+      def preprocess_fn(x):
+        x = tf.expand_dims(x, 0)
+        x = tf.image.resize_bilinear(x, [299, 299],
+                                       align_corners=False)
+        return(tf.squeeze(x, [0]))
+
+      images = tf.map_fn(preprocess_fn, images, dtype=tf.float32)
+
+      net, _ = resnet_v2.resnet_v2_50(
+          images, is_training=False, global_pool=True)
+      output = tf.reshape(net, [shape[0], shape[1], -1])
+      return output
+
+
+class IdentityEmbedder(Embedder):
+  """This embedder just returns the input as the output.
+
+  Used for modalitites that the embedding of the modality is the same as the
+  modality itself. For example, it can be used for one_hot goal.
+  """
+
+  def build(self, images):
+    return images

research/cognitive_planning/envs/configs/active_vision_config.gin

+#-*-Python-*-
+ActiveVisionDatasetEnv.episode_length = 200
+ActiveVisionDatasetEnv.actions = [
+    'right', 'rotate_cw', 'rotate_ccw', 'forward', 'left', 'backward', 'stop'
+]
+ActiveVisionDatasetEnv.confidence_threshold = 0.5
+ActiveVisionDatasetEnv.output_size = 64
+ActiveVisionDatasetEnv.worlds = [
+    'Home_001_1', 'Home_001_2', 'Home_002_1', 'Home_003_1', 'Home_003_2',
+    'Home_004_1', 'Home_004_2', 'Home_005_1', 'Home_005_2', 'Home_006_1',
+    'Home_007_1', 'Home_010_1', 'Home_011_1', 'Home_013_1', 'Home_014_1',
+    'Home_014_2', 'Home_015_1', 'Home_016_1'
+]
+ActiveVisionDatasetEnv.targets = [
+    'tv', 'dining_table', 'fridge', 'microwave', 'couch'
+]
+ActiveVisionDatasetEnv.compute_distance = False
+ActiveVisionDatasetEnv.should_draw_detections = False
+ActiveVisionDatasetEnv.dataset_root = '/usr/local/google/home/kosecka/AVD_Minimal/'
+ActiveVisionDatasetEnv.labelmap_path = 'label_map.txt'
+ActiveVisionDatasetEnv.reward_collision = 0
+ActiveVisionDatasetEnv.reward_goal_range = 2
+ActiveVisionDatasetEnv.num_detection_classes = 90
+ActiveVisionDatasetEnv.segmentation_file_name='sseg_crf'
+ActiveVisionDatasetEnv.detection_folder_name='Detections'
+ActiveVisionDatasetEnv.targets_file_name='annotated_targets'
+ActiveVisionDatasetEnv.shaped_reward=False