我想开发一个分辨率为 256x256 的 DCGAN。为此，我需要使用多个 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.