• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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.eager import context
25from tensorflow.python.framework import config
26from tensorflow.python.framework import ops
27from tensorflow.python.keras import optimizer_v1
28from tensorflow.python.keras.utils import multi_gpu_utils
29from tensorflow.python.keras.utils import np_utils
30from tensorflow.python.platform import test
31
32
33def check_if_compatible_devices(gpus=2):
34  available_devices = [
35      keras.utils.multi_gpu_utils._normalize_device_name(name)
36      for name in keras.utils.multi_gpu_utils._get_available_devices()
37  ]
38  if '/gpu:%d' % (gpus - 1) not in available_devices:
39    return False
40  return True
41
42
43class TestMultiGPUModel(test.TestCase):
44
45  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
46    super(TestMultiGPUModel, self).__init__(methodName)
47    gpu_devices = config.list_physical_devices('GPU')
48    if len(gpu_devices) == 1:
49      # A GPU is available, simulate 2 instead.
50      config.set_logical_device_configuration(gpu_devices[0], [
51          context.LogicalDeviceConfiguration(500),
52          context.LogicalDeviceConfiguration(500)
53      ])
54
55  def test_multi_gpu_test_simple_model(self):
56    gpus = 2
57    num_samples = 1000
58    input_dim = 10
59    output_dim = 1
60    hidden_dim = 10
61    epochs = 2
62    target_gpu_id = [0, 1]
63
64    if not check_if_compatible_devices(gpus=gpus):
65      self.skipTest('multi gpu only')
66
67    with self.cached_session():
68      model = keras.models.Sequential()
69      model.add(keras.layers.Dense(hidden_dim,
70                                   input_shape=(input_dim,)))
71      model.add(keras.layers.Dense(output_dim))
72
73      x = np.random.random((num_samples, input_dim))
74      y = np.random.random((num_samples, output_dim))
75
76      parallel_model = multi_gpu_utils.multi_gpu_model(model, gpus=gpus)
77      parallel_model.compile(loss='mse', optimizer='rmsprop')
78      parallel_model.fit(x, y, epochs=epochs)
79      parallel_model = multi_gpu_utils.multi_gpu_model(
80          model, gpus=target_gpu_id)
81      parallel_model.compile(loss='mse', optimizer='rmsprop')
82      parallel_model.fit(x, y, epochs=epochs)
83
84  def test_multi_gpu_test_multi_io_model(self):
85    gpus = 2
86    num_samples = 1000
87    input_dim_a = 10
88    input_dim_b = 5
89    output_dim_a = 1
90    output_dim_b = 2
91    hidden_dim = 10
92    epochs = 2
93    target_gpu_id = [0, 1]
94
95    if not check_if_compatible_devices(gpus=gpus):
96      self.skipTest('multi gpu only')
97
98    with self.cached_session():
99      input_a = keras.Input((input_dim_a,))
100      input_b = keras.Input((input_dim_b,))
101      a = keras.layers.Dense(hidden_dim)(input_a)
102      b = keras.layers.Dense(hidden_dim)(input_b)
103      c = keras.layers.concatenate([a, b])
104      output_a = keras.layers.Dense(output_dim_a)(c)
105      output_b = keras.layers.Dense(output_dim_b)(c)
106      model = keras.models.Model([input_a, input_b], [output_a, output_b])
107
108      a_x = np.random.random((num_samples, input_dim_a))
109      b_x = np.random.random((num_samples, input_dim_b))
110      a_y = np.random.random((num_samples, output_dim_a))
111      b_y = np.random.random((num_samples, output_dim_b))
112
113      parallel_model = multi_gpu_utils.multi_gpu_model(model, gpus=gpus)
114      parallel_model.compile(loss='mse', optimizer='rmsprop')
115      parallel_model.fit([a_x, b_x], [a_y, b_y], epochs=epochs)
116
117      parallel_model = multi_gpu_utils.multi_gpu_model(
118          model, gpus=target_gpu_id)
119      parallel_model.compile(loss='mse', optimizer='rmsprop')
120      parallel_model.fit([a_x, b_x], [a_y, b_y], epochs=epochs)
121
122  def test_multi_gpu_test_invalid_devices(self):
123    if not check_if_compatible_devices(gpus=2):
124      self.skipTest('multi gpu only')
125
126    with self.cached_session():
127      input_shape = (1000, 10)
128      model = keras.models.Sequential()
129      model.add(keras.layers.Dense(10,
130                                   activation='relu',
131                                   input_shape=input_shape[1:]))
132      model.add(keras.layers.Dense(1, activation='sigmoid'))
133      model.compile(loss='mse', optimizer='rmsprop')
134
135      x = np.random.random(input_shape)
136      y = np.random.random((input_shape[0], 1))
137      with self.assertRaises(ValueError):
138        parallel_model = multi_gpu_utils.multi_gpu_model(
139            model, gpus=len(keras.backend._get_available_gpus()) + 1)
140        parallel_model.fit(x, y, epochs=2)
141
142      with self.assertRaises(ValueError):
143        parallel_model = multi_gpu_utils.multi_gpu_model(
144            model, gpus=[0, 2, 4, 6, 8])
145        parallel_model.fit(x, y, epochs=2)
146
147      with self.assertRaises(ValueError):
148        parallel_model = multi_gpu_utils.multi_gpu_model(model, gpus=1)
149        parallel_model.fit(x, y, epochs=2)
150
151      with self.assertRaises(ValueError):
152        parallel_model = multi_gpu_utils.multi_gpu_model(model, gpus=[0])
153        parallel_model.fit(x, y, epochs=2)
154
155  def test_nested_model_with_tensor_input(self):
156    gpus = 2
157    input_dim = 10
158    shape = (input_dim,)
159    num_samples = 16
160    num_classes = 10
161
162    if not check_if_compatible_devices(gpus=gpus):
163      self.skipTest('multi gpu only')
164
165    with ops.Graph().as_default(), self.cached_session():
166      input_shape = (num_samples,) + shape
167      x_train = np.random.randint(0, 255, input_shape)
168      y_train = np.random.randint(0, num_classes, (input_shape[0],))
169
170      y_train = np_utils.to_categorical(y_train, num_classes)
171
172      x_train = x_train.astype('float32')
173      y_train = y_train.astype('float32')
174
175      dataset = data.Dataset.from_tensor_slices((x_train, y_train))
176      dataset = dataset.repeat()
177      dataset = dataset.batch(4)
178      iterator = data.make_one_shot_iterator(dataset)
179
180      inputs, targets = iterator.get_next()
181
182      input_tensor = keras.layers.Input(tensor=inputs)
183
184      model = keras.models.Sequential()
185      model.add(keras.layers.Dense(3,
186                                   input_shape=(input_dim,)))
187      model.add(keras.layers.Dense(num_classes))
188
189      output = model(input_tensor)
190      outer_model = keras.Model(input_tensor, output)
191      parallel_model = multi_gpu_utils.multi_gpu_model(outer_model, gpus=gpus)
192
193      parallel_model.compile(
194          loss='categorical_crossentropy',
195          optimizer=optimizer_v1.RMSprop(lr=0.0001, decay=1e-6),
196          metrics=['accuracy'],
197          target_tensors=[targets])
198      parallel_model.fit(epochs=1, steps_per_epoch=3)
199
200  def test_multi_gpu_with_multi_input_layers(self):
201    gpus = 2
202
203    if not check_if_compatible_devices(gpus=gpus):
204      self.skipTest('multi gpu only')
205
206    with self.cached_session():
207      inputs = keras.Input((4, 3))
208      init_state = keras.Input((3,))
209      outputs = keras.layers.SimpleRNN(
210          3, return_sequences=True)(inputs, initial_state=init_state)
211      x = [np.random.randn(2, 4, 3), np.random.randn(2, 3)]
212      y = np.random.randn(2, 4, 3)
213      model = keras.Model([inputs, init_state], outputs)
214      parallel_model = multi_gpu_utils.multi_gpu_model(model, gpus=gpus)
215      parallel_model.compile(loss='mean_squared_error', optimizer='adam')
216      parallel_model.train_on_batch(x, y)
217
218  def test_multi_gpu_with_siamese_network(self):
219    gpus = 2
220
221    if not check_if_compatible_devices(gpus=gpus):
222      self.skipTest('multi gpu only')
223
224    with self.cached_session():
225      input_shape = (3,)
226      nested_model = keras.models.Sequential([
227          keras.layers.Dense(32, input_shape=input_shape),
228          keras.layers.Dense(1)
229      ], name='nested')
230
231      input1 = keras.Input(input_shape)
232      input2 = keras.Input(input_shape)
233      score1 = nested_model(input1)
234      score2 = nested_model(input2)
235      score_sum = keras.layers.Add(name='add')([score1, score2])
236
237      siamese = keras.models.Model(inputs=[input1, input2],
238                                   outputs=[score_sum, score1, score2],
239                                   name='siamese')
240      parallel_siamese = multi_gpu_utils.multi_gpu_model(siamese, gpus)
241      self.assertEqual(parallel_siamese.output_names,
242                       ['add', 'nested', 'nested_1'])
243
244if __name__ == '__main__':
245  test.main()
246