def GraphFn(self, x1, x2):
x = x1
q = math_ops.abs(x)
q = q + 1.0
q = gen_math_ops.exp(q)
q = gen_math_ops.log(q)
q = array_ops.squeeze(q, axis=-2)
q = math_ops.abs(q)
q = q + 2.2
q = gen_math_ops.sqrt(q)
q = gen_math_ops.rsqrt(q)
q = math_ops.negative(q)
q = array_ops.squeeze(q, axis=3)
q = math_ops.abs(q)
q = q + 3.0
a = gen_math_ops.reciprocal(q)
# this chain of operations has a batch size of 5, which is different from
# the batch size for the other operations.
x = constant_op.constant(np.random.randn(5, 8, 12), dtype=x.dtype)
q = math_ops.abs(x)
q = q + 2.0
q = gen_math_ops.exp(q)
q = gen_math_ops.log(q)
q = math_ops.abs(q)
q = q + 2.1
q = gen_math_ops.sqrt(q)
q = gen_math_ops.rsqrt(q)
q = math_ops.negative(q)
q = math_ops.abs(q)
q = q + 4.0
b = gen_math_ops.reciprocal(q)
# TODO(jie): this one will break, broadcasting on batch.
x = x2
q = math_ops.abs(x)
q = q + 5.0
q = gen_math_ops.exp(q)
q = array_ops.squeeze(q, axis=[-1, -2, 3])
q = gen_math_ops.log(q)
q = math_ops.abs(q)
q = q + 5.1
q = gen_array_ops.reshape(q, [12, 5, 1, 1, 8, 1, 12])
q = array_ops.squeeze(q, axis=[5, 2, 3])
q = gen_math_ops.sqrt(q)
q = math_ops.abs(q)
q = q + 5.2
q = gen_math_ops.rsqrt(q)
q = math_ops.negative(q)
q = math_ops.abs(q)
q = q + 5.3
c = gen_math_ops.reciprocal(q)
q = a * b
q = q / c
return array_ops.squeeze(q, name="output_0")
def get_updates(self, loss, params):
grads = self.get_gradients(loss, params)
self.updates = []
lr = self.lr
if self.initial_decay > 0:
lr = lr * ( 1. / (1. + self.decay * math_ops.cast(self.iterations,K.dtype(self.decay))) )
with ops.control_dependencies([state_ops.assign_add(self.iterations, 1)]):
t = math_ops.cast(self.iterations, K.floatx())
lr_t = gen_math_ops.sqrt(1. - math_ops.pow(self.beta_2, t)) / (1. - math_ops.pow(self.beta_1, t))
lower_bound = self.lr_boost * (1. - 1. / (self.gamma * t + 1.))
upper_bound = self.lr_boost * (1. + 1. / (self.gamma * t))
if self.sgdcorr:
m_rate = 1. - self.beta_1 / (self.gamma * t + 1.)
else:
m_rate = 1. - self.beta_1
ms = [K.zeros(K.int_shape(p), dtype=K.dtype(p)) for p in params]
vs = [K.zeros(K.int_shape(p), dtype=K.dtype(p)) for p in params]
if self.amsgrad:
vhats = [K.zeros(K.int_shape(p), dtype=K.dtype(p)) for p in params]
else:
vhats = [K.zeros(1) for _ in params]
self.weights = [self.iterations] + ms + vs + vhats
for p, g, m, v, vhat in zip(params, grads, ms, vs, vhats):
m_t = (self.beta_1 * m) + m_rate * g
v_t = (self.beta_2 * v) + (1. - self.beta_2) * math_ops.square(g)
if self.amsgrad:
vhat_t = math_ops.maximum(vhat, v_t)
lr_v = gen_math_ops.reciprocal(gen_math_ops.sqrt(vhat_t) + self.epsilon)
self.updates.append(state_ops.assign(vhat, vhat_t))
else:
lr_v = gen_math_ops.reciprocal(gen_math_ops.sqrt(v_t) + self.epsilon)
lr_bound = gen_math_ops.minimum(gen_math_ops.maximum(lr_t * lr_v, lower_bound), upper_bound)
p_t = p - lr * lr_bound * m_t
self.updates.append(state_ops.assign(m, m_t))
self.updates.append(state_ops.assign(v, v_t))
new_p = p_t
# Apply constraints.
if getattr(p, 'constraint', None) is not None:
new_p = p.constraint(new_p)
self.updates.append(state_ops.assign(p, new_p))
return self.updates
def GraphFn(self, x1, x2):
# Two paths: first with rank 2 input, second with rank 4 input.
outputs = []
xs = [x1, x2]
for i in range(2):
x = xs[i]
c = constant_op.constant(1.0, name="c%d_1" % i)
q = math_ops.add(x, c, name="add%d_1" % i)
q = math_ops.abs(q, name="abs%d_1" % i)
c = constant_op.constant(2.2, name="c%d_2" % i)
q = math_ops.add(q, c, name="add%d_2" % i)
q = math_ops.abs(q, name="abs%d_2" % i)
c = constant_op.constant(3.0, name="c%d_3" % i)
q = math_ops.add(q, c, name="add%d_3" % i)
if i == 0:
axis = constant_op.constant(-1,
dtype=dtypes.int32,
name="axis")
for j in range(2):
q = array_ops.expand_dims(q,
axis,
name="expand%d_%d" % (i, j))
q = self.trt_incompatible_op(q)
q = gen_math_ops.reciprocal(q, name="reciprocal%d" % i)
outputs.append(q)
# Combine both paths
q = math_ops.add(outputs[0], outputs[1], name="add")
return array_ops.squeeze(q, name="output_0")
def GetParams(self):
"""Test for rank 2 input in TF-TRT."""
input_names = ["input", "input2"]
# Two paths: first with rank 2 input, second with rank 4 input.
input_dims = [[12, 5], [12, 5, 2, 2]]
output_name = "output"
g = ops.Graph()
with g.as_default():
outputs = []
for i in range(2):
x = array_ops.placeholder(
dtype=dtypes.float32, shape=input_dims[i], name=input_names[i])
c = constant_op.constant(1.0, name="c%d_1" % i)
q = math_ops.add(x, c, name="add%d_1" % i)
q = math_ops.abs(q, name="abs%d_1" % i)
c = constant_op.constant(2.2, name="c%d_2" % i)
q = math_ops.add(q, c, name="add%d_2" % i)
q = math_ops.abs(q, name="abs%d_2" % i)
c = constant_op.constant(3.0, name="c%d_3" % i)
q = math_ops.add(q, c, name="add%d_3" % i)
if i == 0:
for j in range(2):
q = array_ops.expand_dims(q, -1, name="expand%d_%d" % (i, j))
q = gen_math_ops.reciprocal(q, name="reciprocal%d" % i)
outputs.append(q)
# Combine both paths
q = math_ops.add(outputs[0], outputs[1], name="add")
array_ops.squeeze(q, name=output_name)
return trt_test.TfTrtIntegrationTestParams(
gdef=g.as_graph_def(),
input_names=input_names,
input_dims=input_dims,
output_names=[output_name],
expected_output_dims=[tuple(input_dims[1])])
def _down_scale(self, grads_vars, loss_scale):
# Down scale grads by the loss_scale.
gv = []
inv_loss_scale = gen_math_ops.reciprocal(loss_scale)
for g, v in grads_vars:
if g is not None:
gv.append((g * math_ops.cast(inv_loss_scale, g.dtype.base_dtype), v))
else:
gv.append((g, v))
return gv
def _down_scale(self, grads_vars, loss_scale):
grads_and_vars = []
reciprocal_loss_scale = gen_math_ops.reciprocal(loss_scale)
for grads, variables in grads_vars:
if grads is not None:
grads_and_vars.append((grads * math_ops.cast(
reciprocal_loss_scale, grads.dtype.base_dtype), variables))
else:
grads_and_vars.append((grads, variables))
return grads_and_vars
def GetParams(self):
"""Test for rank 2 input in TF-TRT."""
input_names = ["input", "input2"]
# Two paths: first with rank 2 input, second with rank 4 input.
input_dims = [[12, 5], [12, 5, 2, 2]]
output_name = "output"
g = ops.Graph()
with g.as_default():
outputs = []
for i in range(2):
x = array_ops.placeholder(dtype=dtypes.float32,
shape=input_dims[i],
name=input_names[i])
c = constant_op.constant(1.0, name="c%d_1" % i)
q = math_ops.add(x, c, name="add%d_1" % i)
q = math_ops.abs(q, name="abs%d_1" % i)
c = constant_op.constant(2.2, name="c%d_2" % i)
q = math_ops.add(q, c, name="add%d_2" % i)
q = math_ops.abs(q, name="abs%d_2" % i)
c = constant_op.constant(3.0, name="c%d_3" % i)
q = math_ops.add(q, c, name="add%d_3" % i)
if i == 0:
axis = constant_op.constant(-1,
dtype=dtypes.int32,
name="axis")
for j in range(2):
q = array_ops.expand_dims(q,
axis,
name="expand%d_%d" % (i, j))
q = self.trt_incompatible_op(q)
q = gen_math_ops.reciprocal(q, name="reciprocal%d" % i)
outputs.append(q)
# Combine both paths
q = math_ops.add(outputs[0], outputs[1], name="add")
array_ops.squeeze(q, name=output_name)
return trt_test.TfTrtIntegrationTestParams(
gdef=g.as_graph_def(),
input_names=input_names,
input_dims=[input_dims],
output_names=[output_name],
expected_output_dims=[[input_dims[1]]])
def GetParams(self):
"""Test for unary operations in TF-TRT."""
dtype = dtypes.float32
input_name = "input"
input_dims = [12, 5, 8, 1, 1, 12]
output_name = "output"
input2_name = "input_2"
input2_dims = [12, 5, 8, 1, 12, 1, 1]
g = ops.Graph()
with g.as_default():
x = array_ops.placeholder(dtype=dtype, shape=input_dims, name=input_name)
q = math_ops.abs(x)
q = q + 1.0
q = gen_math_ops.exp(q)
q = gen_math_ops.log(q)
q = array_ops.squeeze(q, axis=-2)
q = math_ops.abs(q)
q = q + 2.2
q = gen_math_ops.sqrt(q)
q = gen_math_ops.rsqrt(q)
q = math_ops.negative(q)
q = array_ops.squeeze(q, axis=3)
q = math_ops.abs(q)
q = q + 3.0
a = gen_math_ops.reciprocal(q)
x = constant_op.constant(np.random.randn(5, 8, 12), dtype=dtype)
q = math_ops.abs(x)
q = q + 2.0
q = gen_math_ops.exp(q)
q = gen_math_ops.log(q)
q = math_ops.abs(q)
q = q + 2.1
q = gen_math_ops.sqrt(q)
q = gen_math_ops.rsqrt(q)
q = math_ops.negative(q)
q = math_ops.abs(q)
q = q + 4.0
b = gen_math_ops.reciprocal(q)
# TODO(jie): this one will break, broadcasting on batch.
x = array_ops.placeholder(
dtype=dtype, shape=input2_dims, name=input2_name)
q = math_ops.abs(x)
q = q + 5.0
q = gen_math_ops.exp(q)
q = array_ops.squeeze(q, axis=[-1, -2, 3])
q = gen_math_ops.log(q)
q = math_ops.abs(q)
q = q + 5.1
q = gen_array_ops.reshape(q, [12, 5, 1, 1, 8, 1, 12])
q = array_ops.squeeze(q, axis=[5, 2, 3])
q = gen_math_ops.sqrt(q)
q = math_ops.abs(q)
q = q + 5.2
q = gen_math_ops.rsqrt(q)
q = math_ops.negative(q)
q = math_ops.abs(q)
q = q + 5.3
c = gen_math_ops.reciprocal(q)
q = a * b
q = q / c
array_ops.squeeze(q, name=output_name)
return trt_test.TfTrtIntegrationTestParams(
gdef=g.as_graph_def(),
input_names=[input_name, input2_name],
input_dims=[input_dims, input2_dims],
output_names=[output_name],
expected_output_dims=[(12, 5, 8, 12)])
def GetParams(self):
"""Test for unary operations in TF-TRT."""
dtype = dtypes.float32
input_name = "input"
input_dims = [12, 5, 8, 1, 1, 12]
input2_name = "input_2"
input2_dims = [12, 5, 8, 1, 12, 1, 1]
g = ops.Graph()
with g.as_default():
x = array_ops.placeholder(dtype=dtype, shape=input_dims, name=input_name)
q = math_ops.abs(x)
q = q + 1.0
q = gen_math_ops.exp(q)
q = gen_math_ops.log(q)
q = array_ops.squeeze(q, axis=-2)
q = math_ops.abs(q)
q = q + 2.2
q = gen_math_ops.sqrt(q)
q = gen_math_ops.rsqrt(q)
q = math_ops.negative(q)
q = array_ops.squeeze(q, axis=3)
q = math_ops.abs(q)
q = q + 3.0
a = gen_math_ops.reciprocal(q)
x = constant_op.constant(np.random.randn(5, 8, 12), dtype=dtype)
q = math_ops.abs(x)
q = q + 2.0
q = gen_math_ops.exp(q)
q = gen_math_ops.log(q)
q = math_ops.abs(q)
q = q + 2.1
q = gen_math_ops.sqrt(q)
q = gen_math_ops.rsqrt(q)
q = math_ops.negative(q)
q = math_ops.abs(q)
q = q + 4.0
b = gen_math_ops.reciprocal(q)
# TODO(jie): this one will break, broadcasting on batch.
x = array_ops.placeholder(
dtype=dtype, shape=input2_dims, name=input2_name)
q = math_ops.abs(x)
q = q + 5.0
q = gen_math_ops.exp(q)
q = array_ops.squeeze(q, axis=[-1, -2, 3])
q = gen_math_ops.log(q)
q = math_ops.abs(q)
q = q + 5.1
q = gen_array_ops.reshape(q, [12, 5, 1, 1, 8, 1, 12])
q = array_ops.squeeze(q, axis=[5, 2, 3])
q = gen_math_ops.sqrt(q)
q = math_ops.abs(q)
q = q + 5.2
q = gen_math_ops.rsqrt(q)
q = math_ops.negative(q)
q = math_ops.abs(q)
q = q + 5.3
c = gen_math_ops.reciprocal(q)
q = a * b
q = q / c
array_ops.squeeze(q, name=self.output_name)
return trt_test.TfTrtIntegrationTestParams(
gdef=g.as_graph_def(),
input_names=[input_name, input2_name],
input_dims=[input_dims, input2_dims],
num_expected_engines=5,
expected_output_dims=(12, 5, 8, 12),
allclose_atol=1.e-03,
allclose_rtol=1.e-03)
def get_updates(self, loss, params):
grads = self.get_gradients(loss, params)
self.updates = [state_ops.assign_add(self.iterations, 1)]
lr = self.lr
if self.initial_decay > 0:
lr = lr * ( 1. / (1. + self.decay * math_ops.cast(self.iterations,K.dtype(self.decay))) )
t = math_ops.cast(self.iterations, K.floatx()) + 1
lower_bound = self.lr_boost * (1. - 1. / (self.gamma * t + 1.))
upper_bound = self.lr_boost * (1. + 1. / (self.gamma * t))
if self.sgdcorr:
m_rate = 1. - self.beta_1 / (self.gamma * t + 1.)
else:
m_rate = 1. - self.beta_1
# Due to the recommendations in [2], i.e. warming momentum schedule
momentum_cache_t = self.beta_1 * (
1. - 0.5 *
(math_ops.pow(K.cast_to_floatx(0.96), t * self.schedule_decay)))
momentum_cache_t_1 = self.beta_1 * (
1. - 0.5 *
(math_ops.pow(K.cast_to_floatx(0.96), (t + 1) * self.schedule_decay)))
m_schedule_new = self.m_schedule * momentum_cache_t
m_schedule_next = self.m_schedule * momentum_cache_t * momentum_cache_t_1
self.updates.append((self.m_schedule, m_schedule_new))
ms = [K.zeros(K.int_shape(p), dtype=K.dtype(p)) for p in params]
vs = [K.zeros(K.int_shape(p), dtype=K.dtype(p)) for p in params]
if self.amsgrad:
vhats = [K.zeros(K.int_shape(p), dtype=K.dtype(p)) for p in params]
else:
vhats = [K.zeros(1) for _ in params]
self.weights = [self.iterations] + ms + vs + vhats
for p, g, m, v, vhat in zip(params, grads, ms, vs, vhats):
# the following equations given in [1]
g_prime = g / (1. - m_schedule_new)
m_t = self.beta_1 * m + m_rate * g
m_t_prime = m_t / (1. - m_schedule_next)
v_t = self.beta_2 * v + (1. - self.beta_2) * math_ops.square(g)
if self.amsgrad:
vhat_t = math_ops.maximum(vhat, v_t)
self.updates.append(state_ops.assign(vhat, vhat_t))
v_t_prime = vhat_t / (1. - math_ops.pow(self.beta_2, t))
else:
v_t_prime = v_t / (1. - math_ops.pow(self.beta_2, t))
m_t_bar = (m_rate / (1.-self.beta_1)) * (1. - momentum_cache_t) * g_prime + momentum_cache_t_1 * m_t_prime
beta_1_reduce = 1. - math_ops.pow(self.beta_1, t)
lr_v = gen_math_ops.reciprocal((gen_math_ops.sqrt(v_t_prime) + self.epsilon) * beta_1_reduce)
self.updates.append(state_ops.assign(m, m_t))
self.updates.append(state_ops.assign(v, v_t))
lr_bound = gen_math_ops.minimum(gen_math_ops.maximum(lr_v, lower_bound), upper_bound)
p_t = p - lr * lr_bound * beta_1_reduce * m_t_bar
new_p = p_t
# Apply constraints.
if getattr(p, 'constraint', None) is not None:
new_p = p.constraint(new_p)
self.updates.append(state_ops.assign(p, new_p))
return self.updates
