tf.distribute.Strategy.scope() 里面必须有什么？

Question

我目前正在使用 tensorflow 2.0 中的分发策略，如下所述https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/distribute/Strategy

我想知道什么必须进入一个with ...scope()块，什么是“可选”。

具体操作如下。我是否必须将 ... 放入 a with ...scope()for distribution 才能工作？：

• 优化器创建
• 数据集创建
• 数据集experimental_distribute_dataset
• apply_gradients 调用
• 循环的数据集迭代
• Experiment_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

您不需要将数据集、数据集迭代循环等放入 内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 的批量大小。您可以阅读更多这里。