1# Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================== 15"""Tests for multi-gpu training utilities.""" 16from __future__ import absolute_import 17from __future__ import division 18from __future__ import print_function 19 20import numpy as np 21 22from tensorflow.python import data 23from tensorflow.python import keras 24from tensorflow.python.platform import test 25 26 27def check_if_compatible_devices(gpus=2): 28 available_devices = [ 29 keras.utils.multi_gpu_utils._normalize_device_name(name) 30 for name in keras.utils.multi_gpu_utils._get_available_devices() 31 ] 32 if '/gpu:%d' % (gpus - 1) not in available_devices: 33 return False 34 return True 35 36 37class TestMultiGPUModel(test.TestCase): 38 39 def test_multi_gpu_test_simple_model(self): 40 gpus = 2 41 num_samples = 1000 42 input_dim = 10 43 output_dim = 1 44 hidden_dim = 10 45 epochs = 2 46 target_gpu_id = [0, 1] 47 48 if not check_if_compatible_devices(gpus=gpus): 49 return 50 51 with self.cached_session(): 52 model = keras.models.Sequential() 53 model.add(keras.layers.Dense(hidden_dim, 54 input_shape=(input_dim,))) 55 model.add(keras.layers.Dense(output_dim)) 56 57 x = np.random.random((num_samples, input_dim)) 58 y = np.random.random((num_samples, output_dim)) 59 60 parallel_model = keras.utils.multi_gpu_model(model, gpus=gpus) 61 parallel_model.compile(loss='mse', optimizer='rmsprop') 62 parallel_model.fit(x, y, epochs=epochs) 63 parallel_model = keras.utils.multi_gpu_model(model, gpus=target_gpu_id) 64 parallel_model.compile(loss='mse', optimizer='rmsprop') 65 parallel_model.fit(x, y, epochs=epochs) 66 67 def test_multi_gpu_test_multi_io_model(self): 68 gpus = 2 69 num_samples = 1000 70 input_dim_a = 10 71 input_dim_b = 5 72 output_dim_a = 1 73 output_dim_b = 2 74 hidden_dim = 10 75 epochs = 2 76 target_gpu_id = [0, 1] 77 78 if not check_if_compatible_devices(gpus=gpus): 79 return 80 81 with self.cached_session(): 82 input_a = keras.Input((input_dim_a,)) 83 input_b = keras.Input((input_dim_b,)) 84 a = keras.layers.Dense(hidden_dim)(input_a) 85 b = keras.layers.Dense(hidden_dim)(input_b) 86 c = keras.layers.concatenate([a, b]) 87 output_a = keras.layers.Dense(output_dim_a)(c) 88 output_b = keras.layers.Dense(output_dim_b)(c) 89 model = keras.models.Model([input_a, input_b], [output_a, output_b]) 90 91 a_x = np.random.random((num_samples, input_dim_a)) 92 b_x = np.random.random((num_samples, input_dim_b)) 93 a_y = np.random.random((num_samples, output_dim_a)) 94 b_y = np.random.random((num_samples, output_dim_b)) 95 96 parallel_model = keras.utils.multi_gpu_model(model, gpus=gpus) 97 parallel_model.compile(loss='mse', optimizer='rmsprop') 98 parallel_model.fit([a_x, b_x], [a_y, b_y], epochs=epochs) 99 100 parallel_model = keras.utils.multi_gpu_model(model, gpus=target_gpu_id) 101 parallel_model.compile(loss='mse', optimizer='rmsprop') 102 parallel_model.fit([a_x, b_x], [a_y, b_y], epochs=epochs) 103 104 def test_multi_gpu_test_invalid_devices(self): 105 if not check_if_compatible_devices(gpus=2): 106 return 107 108 with self.cached_session(): 109 input_shape = (1000, 10) 110 model = keras.models.Sequential() 111 model.add(keras.layers.Dense(10, 112 activation='relu', 113 input_shape=input_shape[1:])) 114 model.add(keras.layers.Dense(1, activation='sigmoid')) 115 model.compile(loss='mse', optimizer='rmsprop') 116 117 x = np.random.random(input_shape) 118 y = np.random.random((input_shape[0], 1)) 119 with self.assertRaises(ValueError): 120 parallel_model = keras.utils.multi_gpu_model( 121 model, gpus=len(keras.backend._get_available_gpus()) + 1) 122 parallel_model.fit(x, y, epochs=2) 123 124 with self.assertRaises(ValueError): 125 parallel_model = keras.utils.multi_gpu_model( 126 model, gpus=[0, 2, 4, 6, 8]) 127 parallel_model.fit(x, y, epochs=2) 128 129 with self.assertRaises(ValueError): 130 parallel_model = keras.utils.multi_gpu_model(model, gpus=1) 131 parallel_model.fit(x, y, epochs=2) 132 133 with self.assertRaises(ValueError): 134 parallel_model = keras.utils.multi_gpu_model(model, gpus=[0]) 135 parallel_model.fit(x, y, epochs=2) 136 137 def test_nested_model_with_tensor_input(self): 138 gpus = 2 139 input_dim = 10 140 shape = (input_dim,) 141 num_samples = 16 142 num_classes = 10 143 144 if not check_if_compatible_devices(gpus=gpus): 145 return 146 147 with self.cached_session(): 148 input_shape = (num_samples,) + shape 149 x_train = np.random.randint(0, 255, input_shape) 150 y_train = np.random.randint(0, num_classes, (input_shape[0],)) 151 152 y_train = keras.utils.to_categorical(y_train, num_classes) 153 154 x_train = x_train.astype('float32') 155 y_train = y_train.astype('float32') 156 157 dataset = data.Dataset.from_tensor_slices((x_train, y_train)) 158 dataset = dataset.repeat() 159 dataset = dataset.batch(4) 160 iterator = data.make_one_shot_iterator(dataset) 161 162 inputs, targets = iterator.get_next() 163 164 input_tensor = keras.layers.Input(tensor=inputs) 165 166 model = keras.models.Sequential() 167 model.add(keras.layers.Dense(3, 168 input_shape=(input_dim,))) 169 model.add(keras.layers.Dense(num_classes)) 170 171 output = model(input_tensor) 172 outer_model = keras.Model(input_tensor, output) 173 parallel_model = keras.utils.multi_gpu_model(outer_model, gpus=gpus) 174 175 parallel_model.compile( 176 loss='categorical_crossentropy', 177 optimizer=keras.optimizers.RMSprop(lr=0.0001, decay=1e-6), 178 metrics=['accuracy'], 179 target_tensors=[targets]) 180 parallel_model.fit(epochs=1, steps_per_epoch=3) 181 182 def test_multi_gpu_with_multi_input_layers(self): 183 gpus = 2 184 185 if not check_if_compatible_devices(gpus=gpus): 186 return 187 188 with self.cached_session(): 189 inputs = keras.Input((4, 3)) 190 init_state = keras.Input((3,)) 191 outputs = keras.layers.SimpleRNN( 192 3, return_sequences=True)(inputs, initial_state=init_state) 193 x = [np.random.randn(2, 4, 3), np.random.randn(2, 3)] 194 y = np.random.randn(2, 4, 3) 195 model = keras.Model([inputs, init_state], outputs) 196 parallel_model = keras.utils.multi_gpu_model(model, gpus=gpus) 197 parallel_model.compile(loss='mean_squared_error', optimizer='adam') 198 parallel_model.train_on_batch(x, y) 199 200 def test_multi_gpu_with_siamese_network(self): 201 gpus = 2 202 203 if not check_if_compatible_devices(gpus=gpus): 204 return 205 206 with self.cached_session(): 207 input_shape = (3,) 208 nested_model = keras.models.Sequential([ 209 keras.layers.Dense(32, input_shape=input_shape), 210 keras.layers.Dense(1) 211 ], name='nested') 212 213 input1 = keras.Input(input_shape) 214 input2 = keras.Input(input_shape) 215 score1 = nested_model(input1) 216 score2 = nested_model(input2) 217 score_sum = keras.layers.Add(name='add')([score1, score2]) 218 219 siamese = keras.models.Model(inputs=[input1, input2], 220 outputs=[score_sum, score1, score2], 221 name='siamese') 222 parallel_siamese = keras.utils.multi_gpu_model(siamese, gpus) 223 self.assertEqual(parallel_siamese.output_names, 224 ['add', 'nested_1', 'nested_2']) 225 226if __name__ == '__main__': 227 test.main() 228