Search code examples
pythontensorflowmathematical-optimization

Where is the code Tensorflow uses to implement RMSProp


I've cloned Tensorflow from https://github.com/tensorflow/tensorflow.git and am trying to find the code where RMS_Prop is implemented.

I have found the file tensorflow/tensorflow/python/training/rmsprop.py, which has calls to training_ops.apply_centered_rms_prop and to training_ops.apply_rms_prop, which seem to be the methods I'm looking for.

I can see that rmsprop.py imports training ops with this import statement:

from tensorflow.python.training import training_ops

but, when I look at training_ops.py, all I see is the following code:

"""Python wrappers for training ops."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from tensorflow.python.training import gen_training_ops  # pylint: disable=unused-import
# go/tf-wildcard-import
# pylint: disable=wildcard-import
from tensorflow.python.training.gen_training_ops import *
# pylint: enable=wildcard-import

I do not see gen_training_ops in the tensorflow/tensorflow/python/training directory and the __init__ file is empty.

What is it I'm not understanding here, and where can I find the actual code that directly implements rmsprop?


Solution

  • The main computational code is implemented in C++; the Python layer abstracts this. The files you are interested in are the CPU kernel and GPU kernel. The kernel for RMSProp starts on lines 500 and 408 respectively.