Universal Multiple Optimizer Wrapper with Layer Assignment

[hyang0129]

I am proposing a universal wrapper for a layer-wise set of distinct optimizers. This will allow each optimizer in a set of optimizers to apply the gradient to its layer-wise variables. This will enable discriminative layer training. It will also enable any training method that applies any combination of optimizers to any combination of layers at any combination of hyperparameters for those optimizers.

The optimizer wrapper will consume a list of instantiated optimizers and layers, referred to as an optimizer spec. For each optimizer spec, it will allocate the correct gradients and variables and then call the apply_gradients method for the optimizer. This allows the optimizer wrapper to leverage all implementations of optimizer specific operations, notably the resource apply methods.

A prototype of the optimizer wrapper is available at this link. link

The prototype works on both TPU and CPU.

I am willing to contribute this code.

This will not change any existing classes. Instead, it will act as a wrapper that allows a model to use multiple optimizers.

Discriminative layer training is most beneficial for fine tuning pretrained models. Users looking to apply existing technology will benefit in reduced training time and improved transfer learning.