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)

Yet another reason to stick optimizer in EstimatorSpec