如何使用Tensorflow 2.0将多个GPU用于DCGAN-RuntimeError:复制本地变量只能在复制上下文中分配

时间:2019-08-13 15:47:32

标签: tensorflow dcgan multiple-gpu

我想开发一个分辨率为256x256的DCGAN。为此,我需要使用多个GPU,因为只有一个GPU是不够的,并且可能会花费太多时间。

我按照此链接上的文档中说明的步骤进行操作 https://www.tensorflow.org/beta/guide/distribute_strategy

在我使用的脚本顶部

strategy = tf.distribute.MirroredStrategy()

然后在我使用的生成器,鉴别器和损失函数中

with strategy.scope():

我得到的错误是:

RuntimeError: Replica-local variables may only be assigned in a replica context.


strategy = tf.distribute.MirroredStrategy()

path = '/my/dataset/path/'
file_paths = [f for f in glob.glob(path + "**/*.jpg", recursive=True)]

tensor_data = np.zeros((len(file_paths), 256, 256, 3)).astype('float32')

for i in range(len(file_paths)): 
  img_tensor = tf.image.decode_image(tf.io.read_file(file_paths[i]))
  tensor_data[i] = img_tensor

for i in range(tensor_data.shape[0]):
  tensor_data[i] = ((tensor_data[i] - 127.5) / 127.5)

BUFFER_SIZE = len(file_paths)
BATCH_SIZE = 256

train_dataset = tf.data.Dataset.from_tensor_slices(tensor_data).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)

def make_generator_model():
    with strategy.scope():
        model = tf.keras.Sequential()
        model.add(layers.Dense(64*64*1536, use_bias=False, input_shape=(100,)))
        model.add(layers.BatchNormalization())
        model.add(layers.LeakyReLU())

        model.add(layers.Reshape((64, 64, 1536)))
        assert model.output_shape == (None, 64, 64, 1536) # Note: None is the batch size

        model.add(layers.Conv2DTranspose(1536, (5, 5), strides=(1, 1), padding='same', use_bias=False))
        assert model.output_shape == (None, 64, 64, 1536)
        model.add(layers.BatchNormalization())
        model.add(layers.LeakyReLU())

        model.add(layers.Conv2DTranspose(768, (5, 5), strides=(2, 2), padding='same', use_bias=False))
        assert model.output_shape == (None, 128, 128, 768)
        model.add(layers.BatchNormalization())
        model.add(layers.LeakyReLU())

        model.add(layers.Conv2DTranspose(3, (5, 5), strides=(2, 2), padding='same', use_bias=False, activation='tanh'))
        assert model.output_shape == (None, 256, 256, 3)

        return model

generator = make_generator_model()

noise = tf.random.normal([1, 100])
generated_image = generator(noise, training=False)

sample = generated_image[0, :, :, :];
sample = tf.cast(sample, tf.int32)

plt.imshow(sample, cmap=None)

def make_discriminator_model():
    with strategy.scope():
        model = tf.keras.Sequential()
        model.add(layers.Conv2D(256, (5, 5), strides=(2, 2), padding='same', input_shape=[256, 256, 3]))
        model.add(layers.LeakyReLU())
        model.add(layers.Dropout(0.3))

        model.add(layers.Conv2D(128, (5, 5), strides=(2, 2), padding='same'))
        model.add(layers.LeakyReLU())
        model.add(layers.Dropout(0.3))

        model.add(layers.Flatten())
        model.add(layers.Dense(1))

    return model

discriminator = make_discriminator_model()
decision = discriminator(generated_image)
print (decision)

cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)

def discriminator_loss(real_output, fake_output):
    with strategy.scope():
        real_loss = cross_entropy(tf.ones_like(real_output), real_output)
        fake_loss = cross_entropy(tf.zeros_like(fake_output), fake_output)
        total_loss = real_loss + fake_loss
        return total_loss

def generator_loss(fake_output):
    with strategy.scope():
        return cross_entropy(tf.ones_like(fake_output), fake_output)

generator_optimizer = tf.keras.optimizers.Adam(1e-4)
discriminator_optimizer = tf.keras.optimizers.Adam(1e-4)

checkpoint_dir = './training_checkpoints/'
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
checkpoint = tf.train.Checkpoint(generator_optimizer=generator_optimizer,
                                 discriminator_optimizer=discriminator_optimizer,
                                 generator=generator,
                                 discriminator=discriminator)

EPOCHS = 2000
noise_dim = 100
num_examples_to_generate = 16

seed = tf.random.normal([num_examples_to_generate, noise_dim])

# Notice the use of `tf.function`
# This annotation causes the function to be "compiled".
@tf.function
def train_step(images):
    noise = tf.random.normal([BATCH_SIZE, noise_dim])

    with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
        generated_images = generator(noise, training=True)

        real_output = discriminator(images, training=True)
        fake_output = discriminator(generated_images, training=True)

        gen_loss = generator_loss(fake_output)
        disc_loss = discriminator_loss(real_output, fake_output)

    gradients_of_generator = gen_tape.gradient(gen_loss, generator.trainable_variables)
    gradients_of_discriminator = disc_tape.gradient(disc_loss, discriminator.trainable_variables)

    generator_optimizer.apply_gradients(zip(gradients_of_generator, generator.trainable_variables))
    discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, discriminator.trainable_variables))

def train(dataset, epochs):
  for epoch in range(epochs):
    start = time.time()

    for image_batch in dataset:
      train_step(image_batch)

    # Produce images for the GIF as we go
    display.clear_output(wait=True)
    generate_and_save_images(generator,
                             epoch + 1,
                             seed)

    # Save the model every 15 epochs
    os.makedirs(os.path.dirname(checkpoint_prefix), exist_ok=True)

    if (epoch + 1) % 50 == 0:
      checkpoint.save(file_prefix = checkpoint_prefix)

    print ('Time for epoch {} is {} sec'.format(epoch + 1, time.time()-start))

  # Generate after the final epoch
  display.clear_output(wait=True)
  generate_and_save_images(generator,epochs,seed)

def generate_and_save_images(model, epoch, test_input):
    # Notice `training` is set to False.
    # This is so all layers run in inference mode (batchnorm).
    predictions = model(test_input, training=False)

    fig = plt.figure(figsize=(4,4))

    for i in range(predictions.shape[0]):
      plt.subplot(8, 8, i+1)
      sample = predictions[i, :, :, :] * 127.5 + 127.5
      sample = tf.cast(sample, tf.int32)
      plt.imshow(sample, cmap=None)
      plt.axis('off')

    filename = './screens/eye-256x256/1/image_at_epoch_{:04d}.png'
    os.makedirs(os.path.dirname(filename), exist_ok=True)
    if (epoch + 1) % 10 == 0:
        plt.savefig(filename.format(epoch))
        plt.show()

get_ipython().run_cell_magic('time', '', 'train(train_dataset, EPOCHS)')

错误如下

Executing op ExperimentalRebatchDataset in device /job:localhost/replica:0/task:0/device:CPU:0
Executing op ExperimentalAutoShardDataset in device /job:localhost/replica:0/task:0/device:CPU:0
Executing op OptimizeDataset in device /job:localhost/replica:0/task:0/device:CPU:0
Executing op ModelDataset in device /job:localhost/replica:0/task:0/device:CPU:0
Executing op MultiDeviceIterator in device /job:localhost/replica:0/task:0/device:CPU:0
Executing op MultiDeviceIteratorInit in device /job:localhost/replica:0/task:0/device:CPU:0
Executing op MultiDeviceIteratorToStringHandle in device /job:localhost/replica:0/task:0/device:CPU:0
Executing op GeneratorDataset in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op GeneratorDataset in device /job:localhost/replica:0/task:0/device:GPU:1
Executing op PrefetchDataset in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op AnonymousIteratorV2 in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op MakeIterator in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op PrefetchDataset in device /job:localhost/replica:0/task:0/device:GPU:1
Executing op AnonymousIteratorV2 in device /job:localhost/replica:0/task:0/device:GPU:1
Executing op MakeIterator in device /job:localhost/replica:0/task:0/device:GPU:1
Executing op IteratorGetNextSync in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op IteratorGetNextSync in device /job:localhost/replica:0/task:0/device:GPU:1
Executing op DestroyResourceOp in device /job:localhost/replica:0/task:0/device:CPU:0
Executing op DeleteIterator in device /job:localhost/replica:0/task:0/device:GPU:1
Executing op DeleteIterator in device /job:localhost/replica:0/task:0/device:GPU:0
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<timed exec> in <module>

<ipython-input-20-88a9879432c7> in train(dataset, epochs)
      4 
      5     for image_batch in dataset:
----> 6       train_step(image_batch)
      7 
      8     # Produce images for the GIF as we go

/usr/local/lib/python3.5/dist-packages/tensorflow/python/eager/def_function.py in __call__(self, *args, **kwds)
    414     # This is the first call of __call__, so we have to initialize.
    415     initializer_map = {}
--> 416     self._initialize(args, kwds, add_initializers_to=initializer_map)
    417     if self._created_variables:
    418       try:

/usr/local/lib/python3.5/dist-packages/tensorflow/python/eager/def_function.py in _initialize(self, args, kwds, add_initializers_to)
    357     self._concrete_stateful_fn = (
    358         self._stateful_fn._get_concrete_function_internal_garbage_collected(  # pylint: disable=protected-access
--> 359             *args, **kwds))
    360 
    361     def invalid_creator_scope(*unused_args, **unused_kwds):

/usr/local/lib/python3.5/dist-packages/tensorflow/python/eager/function.py in _get_concrete_function_internal_garbage_collected(self, *args, **kwargs)
   1358     if self.input_signature:
   1359       args, kwargs = None, None
-> 1360     graph_function, _, _ = self._maybe_define_function(args, kwargs)
   1361     return graph_function
   1362 

/usr/local/lib/python3.5/dist-packages/tensorflow/python/eager/function.py in _maybe_define_function(self, args, kwargs)
   1646       graph_function = self._function_cache.primary.get(cache_key, None)
   1647       if graph_function is None:
-> 1648         graph_function = self._create_graph_function(args, kwargs)
   1649         self._function_cache.primary[cache_key] = graph_function
   1650       return graph_function, args, kwargs

/usr/local/lib/python3.5/dist-packages/tensorflow/python/eager/function.py in _create_graph_function(self, args, kwargs, override_flat_arg_shapes)
   1539             arg_names=arg_names,
   1540             override_flat_arg_shapes=override_flat_arg_shapes,
-> 1541             capture_by_value=self._capture_by_value),
   1542         self._function_attributes)
   1543 

/usr/local/lib/python3.5/dist-packages/tensorflow/python/framework/func_graph.py in func_graph_from_py_func(name, python_func, args, kwargs, signature, func_graph, autograph, autograph_options, add_control_dependencies, arg_names, op_return_value, collections, capture_by_value, override_flat_arg_shapes)
    714                                           converted_func)
    715 
--> 716       func_outputs = python_func(*func_args, **func_kwargs)
    717 
    718       # invariant: `func_outputs` contains only Tensors, CompositeTensors,

/usr/local/lib/python3.5/dist-packages/tensorflow/python/eager/def_function.py in wrapped_fn(*args, **kwds)
    307         # __wrapped__ allows AutoGraph to swap in a converted function. We give
    308         # the function a weak reference to itself to avoid a reference cycle.
--> 309         return weak_wrapped_fn().__wrapped__(*args, **kwds)
    310     weak_wrapped_fn = weakref.ref(wrapped_fn)
    311 

/usr/local/lib/python3.5/dist-packages/tensorflow/python/framework/func_graph.py in wrapper(*args, **kwargs)
    704           except Exception as e:  # pylint:disable=broad-except
    705             if hasattr(e, "ag_error_metadata"):
--> 706               raise e.ag_error_metadata.to_exception(type(e))
    707             else:
    708               raise

RuntimeError: in converted code:

    <ipython-input-19-d2ffe8a85706>:9 train_step  *
        generated_images = generator(noise, training=True)
    /usr/local/lib/python3.5/dist-packages/tensorflow/python/keras/engine/base_layer.py:667 __call__
        outputs = call_fn(inputs, *args, **kwargs)
    /usr/local/lib/python3.5/dist-packages/tensorflow/python/keras/engine/sequential.py:248 call
        return super(Sequential, self).call(inputs, training=training, mask=mask)
    /usr/local/lib/python3.5/dist-packages/tensorflow/python/keras/engine/network.py:753 call
        return self._run_internal_graph(inputs, training=training, mask=mask)
    /usr/local/lib/python3.5/dist-packages/tensorflow/python/keras/engine/network.py:895 _run_internal_graph
        output_tensors = layer(computed_tensors, **kwargs)
    /usr/local/lib/python3.5/dist-packages/tensorflow/python/keras/engine/base_layer.py:667 __call__
        outputs = call_fn(inputs, *args, **kwargs)
    /usr/local/lib/python3.5/dist-packages/tensorflow/python/keras/layers/normalization.py:782 call
        self.add_update(mean_update)
    /usr/local/lib/python3.5/dist-packages/tensorflow/python/util/deprecation.py:507 new_func
        return func(*args, **kwargs)
    /usr/local/lib/python3.5/dist-packages/tensorflow/python/keras/engine/base_layer.py:1095 add_update
        update()
    /usr/local/lib/python3.5/dist-packages/tensorflow/python/keras/layers/normalization.py:775 mean_update
        return tf_utils.smart_cond(training, true_branch, false_branch)
    /usr/local/lib/python3.5/dist-packages/tensorflow/python/keras/utils/tf_utils.py:58 smart_cond
        pred, true_fn=true_fn, false_fn=false_fn, name=name)
    /usr/local/lib/python3.5/dist-packages/tensorflow/python/framework/smart_cond.py:54 smart_cond
        return true_fn()
    /usr/local/lib/python3.5/dist-packages/tensorflow/python/keras/layers/normalization.py:773 <lambda>
        true_branch = lambda: _do_update(self.moving_mean, new_mean)
    /usr/local/lib/python3.5/dist-packages/tensorflow/python/keras/layers/normalization.py:769 _do_update
        inputs_size)
    /usr/local/lib/python3.5/dist-packages/tensorflow/python/keras/layers/normalization.py:458 _assign_moving_average
        return state_ops.assign_sub(variable, update_delta, name=scope)
    /usr/local/lib/python3.5/dist-packages/tensorflow/python/ops/state_ops.py:164 assign_sub
        return ref.assign_sub(value)
    /usr/local/lib/python3.5/dist-packages/tensorflow/python/distribute/values.py:1394 assign_sub
        _assert_replica_context(self._distribute_strategy)
    /usr/local/lib/python3.5/dist-packages/tensorflow/python/distribute/values.py:1381 _assert_replica_context
        "Replica-local variables may only be assigned in a replica context.")

    RuntimeError: Replica-local variables may only be assigned in a replica context.

1 个答案:

答案 0 :(得分:0)

您需要同时分发数据集,有关详细信息,请参考此URL。  -https://www.tensorflow.org/beta/tutorials/distribute/training_loops

train_dist_dataset = strategy.experimental_distribute_dataset(train_dataset)

模型的每个部分都会在战略范围内(例如优化程序)创建。

with strategy.scope():
  optimizer = tf.keras.optimizers.Adam()