When create an optimizer in Tensorflow, how to deal with AttributeError happens?

Question

I try to incorporate a self-designed optimization algorithm PSGLD into TensorFlow. And that algorithm is similar to the concept of RMSProp. So I didn't create a new Op, but complement PSGLD following RMSProp. My procedure of incorporating is as follows:

1. In Python side, create a psgld.py under the folder of tensorflow\python\training,which represents the Python wrapper. And in psgld.py, define the class of PSGLDOptimizer.

   class PSGLDOptimizer(optimizer.Optimizer)

2. Then, in tensorflow\python\training\training_ops.py, define the shape function of _ApplyPSGLDShape and _SparseApplyPSGLD, for dense and sparse circumstances respectively.

3. For C++ side, in tensorflow\core\ops\training_ops.cc, define the input, output and attribute of ApplyPSGLD Op:
   REGISTER_OP("ApplyPSGLD") .Input("var: Ref(T)") .Input("ms: Ref(T)") .Input("mom: Ref(T)") .Input("lr: T") .Input("decay: T") .Input("epsilon: T") .Input("grad: T") .Output("out: Ref(T)") .Attr("T: numbertype") .Attr("use_locking: bool = false")

4. Meanwhile, also define ApplyPSGLD in the header file of tensorflow\core\kernels\training_ops.h

   template <typename Device, typename T> struct ApplyPSGLD { ... };

5. To realize the computation of our algorithm on C++ side, complement corresponding code in the kernel of tensorflow\core\kernels\training_ops.cc.

After all, when I run tensorflow/models/image/mnist/convolutional.py, and the optimizer is adjusted,

optimizer = tf.train.PSGLDOptimizer(learning_rate).minimize(loss, global_step=batch)

an AttributeError happens: AttributeError: 'module' object has no attribute 'PSGLDOptimizer'

And the environment is TF-0.9, cudnn5. So I ask if someone can give me any advice on this issue or the whole procedure of adding an optimizer.

Answer 1

(I'm assuming that you've rebuilt TensorFlow from source, as Olivier suggested in his comment, and you are trying to construct your optimizer as optimizer = tf.train.PSGLDOptimizer(...).)

To add a symbol to the tf.train namespace, you have to do the following:

1. Add an explicit import to the file tensorflow/python/training/training.py. In that file, you can see imports for, e.g., the tf.train.RMSPropOptimizer class.

2. Either:

   1. Add documentation for your new class, and add @@PSGLDOptimizer to the module docstring. The corresponding line for tf.train.RMSPropOptimizer is here. This marks the class as a publicly documented API symbol.

   2. Add an exception to the whitelist of symbols that are added to __all__ in that file. For example, this line whitelists the symbol tf.train.LooperThread.

For most TensorFlow modules*, the rule is that a symbol can appear in __all__ if it is either (i) publicly documented, or (ii) explicitly whitelisted. If neither condition holds, it will not be accessible through a tf.* name. This is intended to keep the API surface small, and avoid exposing private implementation details that might change between versions.

* Note however that this is a work in progress. At present, a method is considered to be stable only if it is documented in the public API docs.