tensorflow / addons

Useful extra functionality for TensorFlow 2.x maintained by SIG-addons
Apache License 2.0
1.69k stars 611 forks source link

tfa.optimizers.MovingAverage and default strategy #2618

Open plooney opened 2 years ago

plooney commented 2 years ago

I have two machines one with multiple GPUs and one with a single GPU.

I do something like

strategy = tf.distribute.MirroredStrategy() if multi_gpu else tf.distribute.get_strategy()

then in training I do

with strategy.scope(): optimizer.swap_weights()

This works fine on the multi_gpu machine but on the single GPU machine it produces the error

ValueError: Swapping weights must occur under a tf.distribute.Strategy

Is there a way to use this on a single GPU?

fsx950223 commented 2 years ago

MirroredStrategy supports single GPU too.