multi_gpu_estimator_mixin_ex.py

class MyEstimator(TowerEstimator, DNNClassifier):
  pass


def my_optimizer_fn(grads_and_vars, params):
  optimizer = AdagradOptimizer(learning_rate=params.learning_rate, momentum=params.momentum)
  return optimizer.apply_gradients(grads_and_vars)


estimator = MyEstimator(feature_columns=feature_columns, 
                        layers=[100, 80, 60, 40],
                        optimizer_fn=my_optimizer_fn,
                        params=HParams(learning_rate=0.5, momentum=0.1))
  
estimator.fit(input_fn=my_input_fn)

grads_and_vars were computed with tf.gradients, right?

Not with some optimizer.

Correct. They are the result of averaging the gradients across all the towers. So if you want to clip or gate you'll have to do it manually here.

Yet another reason to stick optimizer in EstimatorSpec