def test_LSTM_runtime_with_cond(self): # This test is to demonstrate the graph rewrite of grappler plugin under # the condition that the function returns different number of internal # states. layer = rnn.LSTM(self.rnn_state_size, return_runtime=True) inputs = keras.layers.Input(shape=[self.timestep, self.input_shape], dtype=dtypes.float32) zeros = array_ops.zeros([self.batch, self.output_shape]) dummy_runtime = rnn._runtime(rnn._RUNTIME_UNKNOWN) a = constant_op.constant(0) b = constant_op.constant(1) # Will always run the lstm layer. outputs, runtime = control_flow_ops.cond( gen_math_ops.less(a, b), lambda: layer(inputs), lambda: (zeros, dummy_runtime)) # Expand the runtime so that it is a 1D tensor instead of scalar. # TF model does not work with scalar model output, specially during # aggregation. runtime = keras.layers.Lambda( lambda x: array_ops.expand_dims(x, axis=-1))(runtime) model = keras.models.Model(inputs=inputs, outputs=[outputs, runtime]) self._test_runtime_with_model(model)
def test_unifiedLSTM_with_cond(self): # This test is to demonstrate the graph rewrite of grappler plugin under # the condition that the function returns different number of internal # states. input_shape = 10 rnn_state_size = 8 output_shape = 8 timestep = 4 batch = 100 epoch = 1 with self.cached_session(config=_config, use_gpu=True) as sess: (x_train, y_train), _ = testing_utils.get_test_data( train_samples=batch, test_samples=0, input_shape=(timestep, input_shape), num_classes=output_shape) y_train = keras.utils.to_categorical(y_train, output_shape) layer = rnn.LSTM(rnn_state_size, return_runtime=True) inputs = array_ops.placeholder(dtypes.float32, shape=(None, timestep, input_shape), name='inputs') predict = array_ops.placeholder(dtypes.float32, shape=(None, output_shape), name='predict') zeros = array_ops.zeros([batch, output_shape]) dummy_runtime = rnn._runtime(rnn._RUNTIME_UNKNOWN) a = constant_op.constant(0) b = constant_op.constant(1) # Will always run the lstm layer. outputs, runtime = control_flow_ops.cond( gen_math_ops.less(a, b), lambda: layer(inputs), lambda: (zeros, dummy_runtime)) loss = losses.softmax_cross_entropy(predict, outputs) optimizer = gradient_descent.GradientDescentOptimizer(0.001) train_op = optimizer.minimize(loss) sess.run([variables.global_variables_initializer()]) existing_loss = 0 for _ in range(epoch): loss_value, _, runtime_value = sess.run( [loss, train_op, runtime], { inputs: x_train, predict: y_train }) if test.is_gpu_available(): self.assertEqual(runtime_value, rnn._RUNTIME_GPU) else: self.assertEqual(runtime_value, rnn._RUNTIME_CPU) # Make sure the loss is updated for every epoch # (layer weights properly updated). self.assertNotEqual(existing_loss, loss_value) existing_loss = loss_value
def test_unifiedLSTM_with_cond(self): # This test is to demonstrate the graph rewrite of grappler plugin under # the condition that the function returns different number of internal # states. input_shape = 10 rnn_state_size = 8 output_shape = 8 timestep = 4 batch = 100 epoch = 1 with self.cached_session(config=_config, use_gpu=True) as sess: (x_train, y_train), _ = testing_utils.get_test_data( train_samples=batch, test_samples=0, input_shape=(timestep, input_shape), num_classes=output_shape) y_train = keras.utils.to_categorical(y_train, output_shape) layer = rnn.LSTM(rnn_state_size, return_runtime=True) inputs = array_ops.placeholder( dtypes.float32, shape=(None, timestep, input_shape), name='inputs') predict = array_ops.placeholder( dtypes.float32, shape=(None, output_shape), name='predict') zeros = array_ops.zeros([batch, output_shape]) dummy_runtime = rnn._runtime(rnn._RUNTIME_UNKNOWN) a = constant_op.constant(0) b = constant_op.constant(1) # Will always run the lstm layer. outputs, runtime = control_flow_ops.cond( gen_math_ops.less(a, b), lambda: layer(inputs), lambda: (zeros, dummy_runtime)) loss = losses.softmax_cross_entropy(predict, outputs) optimizer = gradient_descent.GradientDescentOptimizer(0.001) train_op = optimizer.minimize(loss) sess.run([variables.global_variables_initializer()]) existing_loss = 0 for _ in range(epoch): loss_value, _, runtime_value = sess.run([loss, train_op, runtime], { inputs: x_train, predict: y_train }) if test.is_gpu_available(): self.assertEqual(runtime_value, rnn._RUNTIME_GPU) else: self.assertEqual(runtime_value, rnn._RUNTIME_CPU) # Make sure the loss is updated for every epoch # (layer weights properly updated). self.assertNotEqual(existing_loss, loss_value) existing_loss = loss_value