def _time_performance_run_cudnn_lstm(self, test_config, x_train, y_train):
# Get the performance number for standard Cudnn LSTM
input_shape = test_config['input_shape']
rnn_state_size = test_config['rnn_state_size']
timestep = test_config['timestep']
epoch = test_config['epoch']
warmup_epoch = test_config['warmup_epoch']
ops.reset_default_graph()
with self.test_session(use_gpu=True):
cudnn_lstm_layer = CuDNNLSTM(rnn_state_size)
inputs = keras.layers.Input(shape=[timestep, input_shape],
dtype=dtypes.float32)
outputs = cudnn_lstm_layer(inputs)
model = keras.models.Model(inputs, outputs)
model.compile('sgd', 'mse')
total_duration = 0
for i in range(epoch):
start_time = time.time()
model.fit(x_train, y_train)
end_time = time.time()
if i >= warmup_epoch:
duration_per_epoch = end_time - start_time
total_duration += duration_per_epoch
logging.vlog(2, '%s: Time consumed for epoch %d is: %s',
'CuDNN LSTM', i, duration_per_epoch)
logging.info('Average performance for %s per epoch is: %s',
'CuDNN LSTM', (total_duration / epoch))
return total_duration / epoch
def _time_performance_run_cudnn_lstm(self, test_config, x_train, y_train):
# Get the performance number for standard Cudnn LSTM
input_shape = test_config['input_shape']
rnn_state_size = test_config['rnn_state_size']
timestep = test_config['timestep']
cudnn_lstm_layer = CuDNNLSTM(rnn_state_size)
inputs = keras.layers.Input(shape=[timestep, input_shape],
dtype=dtypes.float32)
outputs = cudnn_lstm_layer(inputs)
model = keras.models.Model(inputs, outputs)
model.compile('sgd', 'mse')
sec_per_epoch = self._measure_performance(test_config, model, x_train,
y_train)
logging.info('Average performance for %s per epoch is: %s',
'CuDNN LSTM', sec_per_epoch)
return sec_per_epoch
#
# tf.config.experimental.set_virtual_device_configuration(gpus[0], [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=4.5*1024)])
import numpy as np
from tensorflow.keras.layers import Dense
from tensorflow.python.keras.layers.cudnn_recurrent import CuDNNLSTM
from tensorflow.keras import Sequential
num_samples = 1000000
timesteps = 30
timestep_size = 4
x = np.random.rand(num_samples, timesteps, timestep_size)
y = np.random.rand(num_samples, 1)
print('x shape: {}'.format(x.shape))
print('y shape: {}'.format(y.shape))
# print('Press enter to start training...', flush=True)
# input()
model = Sequential()
model.add(CuDNNLSTM(128, return_sequences=True, input_shape=x.shape[1:]))
model.add(CuDNNLSTM(128, return_sequences=True))
model.add(CuDNNLSTM(64))
model.add(Dense(64, activation='relu'))
model.add(Dense(128, activation='relu'))
model.add(Dense(64, activation='relu'))
model.add(Dense(1))
model.compile(optimizer='adam', loss='mse')
model.fit(x, y, batch_size=16, epochs=500)