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