What has to be inside tf.distribute.Strategy.scope()?

user_7128838 0
Viewed 8226

I am currently playing around with Distribution Strategies in tensorflow 2.0 as described here https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/distribute/Strategy

I am wondering what has to go inside a with ...scope() block and what is "optional".

Specifically the following operations. Do I have to put ... inside a with ...scope() for distribution to work?:

  • Optimizer creation
  • Dataset creation
  • Dataset experimental_distribute_dataset
  • apply_gradients call
  • Dataset iteration for loop
  • experimental_run_v2

I have toyed around a little and my code seems to work even when I use no with ...scope at all. I am confused if this has some side-effects I am just not seeing right now.

Code without 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

Code with 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)

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

Edit: It seems that strat.experimental_run_v2 automatically enters the scope of strat. So why does with strat.scope() exist?

tensorflow2.0 tensorflow python
4 Answers

According to my experiments, the only thing that needs to be declared inside is model creation. If you use use Keras .fit() instead of custom training then model.compile() has to be inside as well.

You can do something like this:

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

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

If you use tf.train.Checkpoint then make sure both its instantiation and the call of checkpoint.resume() are inside the scope.

user_21312000

You do not need to put Dataset, dataset iteration loop etc. inside scope(). You just need to define your sequential model and its compilation inside it. So something like this-

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'])

What it will do is, it will create a replica of model and its parameters on each GPU which will be trained during training. The batch size you will define will be divided by the number of GPUs available and those batches will be sent to those GPUs, for example, if you have batch_size = 64 and you have two GPUs then each GPU will get batch sizes of 32. You can read read more here.

user_88426940
user_38671330

My experience with TPU on Kaggle has shown those effects:

  1. create strategy and connect to TPU before building datasets (tf.dataset.Dataset). Otherwise, especially with more complex dataset loading/preprocessing TPU unit was complaining.

  2. Inside with strategy.scope(): put:

    a) model creation
    b) instantiation of the metrics
    c) compilation of the model

With simple metrics it works even when compiled outside the scope. But when added list of custom metrics, again TPU started complaining and putting all under the .scope() solved the issue.

user_69104710
Related