Question

我的工作量在两个GPU上进行了分区（又称模型分区）。默认情况下，TF / Keras在GPU0上分配所有渐变，但是我想使用colocate_gradients_with_ops在两个GPU之间分配分配。

我正在寻找一种在Keras中实现这一目标的简单方法。我的方法是创建一个从tf.train.AdamOptimizer子类化的新优化器，以将colocate_gradients_with_ops的默认值翻转（从False到True）。另外，我必须用两种方法翻转它！

我正在寻找一种比Keras下面的方法更短，更直接的方法。

class MyAdamOptimizer(tf.train.AdamOptimizer):
    def compute_gradients(self,
                          loss,
                          var_list=None,
                          gate_gradients=tf.train.Optimizer.GATE_OP,
                          aggregation_method=None,
                          colocate_gradients_with_ops=True,
                          grad_loss=None):
        return super(MyAdamOptimizer, self).compute_gradients(
            loss,
            var_list=None,
            gate_gradients=tf.train.Optimizer.GATE_OP,
            aggregation_method=None,
            colocate_gradients_with_ops=True,
            grad_loss=None)

    def minimize(
            loss,
            global_step=None,
            var_list=None,
            gate_gradients=tf.train.Optimizer.GATE_OP,
            aggregation_method=None,
            colocate_gradients_with_ops=True,
            name=None,
            grad_loss=None):
        return super(MyAdamOptimizer, self).minimize(
            loss,
            global_step=None,
            var_list=None,
            gate_gradients=tf.train.Optimizer.GATE_OP,
            aggregation_method=None,
            colocate_gradients_with_ops=True,
            name=None,
            grad_loss=None)

然后我打电话

model.compile(optimizer=MyAdamOptimizer(learning_rate=0.001),
              loss='categorical_crossentropy',
              metrics=['accuracy'])

Answer 1

没有更简单的方法。 Keras AdamOptimizer使用其自己的基本运算符实现。您必须为colocate_gradients_with_ops使用自定义优化器。如果要提高多GPU的性能，可以尝试Keras-MXNet的AdamOptimizer，我们超载了Keras的Optimizer类，并在多GPU上具有更高的效率。您不必更改培训代码。

使用colocate_gradients_with_ops = True在Keras中自定义优化器的正确方法是什么？

1 个答案: