tf.distribute.Strategy.scope()中必须包含什么？

Question

我目前正在使用tensorflow 2.0中的分发策略，如本文所述，docs/python/tf/分发/策略

我想知道在with ...scope()块中必须进行什么，以及什么是“可选的”。

特别是以下操作。我一定要把..。在with ...scope()中运行发行版？：

• 优化器创建
• 数据集创建
• 数据集experimental_distribute_dataset
• apply_gradients调用
• 数据集循环迭代
• experimental_run_v2

我已经玩了一会儿，我的代码似乎工作，即使我根本不使用with ...scope。我很困惑这是否有一些副作用，我只是现在没有看到。

没有scope的代码

strat = tf.distribute.MirroredStrategy()

BATCH_SIZE_PER_REPLICA = 5

print('Replicas: ', strat.num_replicas_in_sync)

global_batch_size = (BATCH_SIZE_PER_REPLICA * strat.num_replicas_in_sync)

dataset = tf.data.Dataset.from_tensors(tf.random.normal([100])).repeat(1000).batch(
    global_batch_size)

g = Model('m', 10, 10, 1, 3)

dist_dataset = strat.experimental_distribute_dataset(dataset)

@tf.function
def train_step(dist_inputs):
  def step_fn(inputs):
    print([(v.name, v.device) for v in g.trainable_variables])
    return g(inputs)

  out = strat.experimental_run_v2(step_fn, args=(dist_inputs,))

for inputs in dist_dataset:
    train_step(inputs)
    break

具有作用域的代码：

strat = tf.distribute.MirroredStrategy()

BATCH_SIZE_PER_REPLICA = 5

print('Replicas: ', strat.num_replicas_in_sync)

global_batch_size = (BATCH_SIZE_PER_REPLICA * strat.num_replicas_in_sync)

with strat.scope():
    dataset = tf.data.Dataset.from_tensors(tf.random.normal([100])).repeat(1000).batch(
        global_batch_size)

    g = Model('m', 10, 10, 1, 3)

    dist_dataset = strat.experimental_distribute_dataset(dataset)

    @tf.function
    def train_step(dist_inputs):
        def step_fn(inputs):
            print([(v.name, v.device) for v in g.trainable_variables])
            return g(inputs)

        out = strat.experimental_run_v2(step_fn, args=(dist_inputs,))

    for inputs in dist_dataset:
        train_step(inputs)
        break

编辑：strat.experimental_run_v2似乎自动进入了strat的范围。那么为什么with strat.scope()会存在呢？

Answer 1

根据我的实验，唯一需要在里面声明的是模型的创建。如果您使用Keras .fit()而不是自定义培训，那么model.compile()也必须在内部。

你可以这样做：

def create_model():
    """ This can be outside of the scope
    """
    ...
    return model

with strategy.scope():
    model = create_model()

如果使用tf.train.Checkpoint，那么确保它的实例化和checkpoint.resume()调用都在作用域中。

Answer 2

您不需要将Dataset、dataset迭代循环等放在scope()中。您只需要定义您的顺序模型及其内部的编译。所以就像这样-

mirrored_strategy = tf.distribute.MirroredStrategy()
with mirrored_strategy.scope():
  model = tf.keras.Sequential()
  model.add(tf.keras.layers.Embedding(vocab_size, 64))
  model.add(tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(64, activation= 'tanh', recurrent_activation= 'sigmoid', recurrent_dropout = 0, unroll = False, use_bias= True)))
  # One or more dense layers.
  # Edit the list in the `for` line to experiment with layer sizes.
  for units in [64, 64]:
    model.add(tf.keras.layers.Dense(units, activation='relu'))
  # Output layer. The first argument is the number of labels.
  model.add(tf.keras.layers.Dense(3, activation='softmax'))
  model.compile(optimizer='adam',
                loss='sparse_categorical_crossentropy',
                metrics=['accuracy'])

它所要做的是，它将在每个GPU上创建一个模型的复制品及其参数，在培训期间进行培训。您将定义的批处理大小将除以可用的GPU数量，这些批将被发送到这些GPU，例如，如果您有batch_size = 64，并且您有两个GPU，那么每个GPU将得到32的批处理大小。您可以阅读更多的这里。

Answer 3

你不需要打电话给strat.scope()。

experimental_run_v2是将您的计算放在strat.scope()中的一种简单方法。

请参阅下面的experimental_run_v2源代码，它实际上将您的fn包装在您的作用域中。

lib.py#L729