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