[Training] Add Adagrad optimizer operator

[wschin]

PR #2314 is a single place for reviewing the whole training story.

This PR proposes a stateless optimizer based on a stochastic gradient method. This operator covers Tensorflow's and Pytorch's ADAGRAD optimizer.

Optimizer itself can be defined as a stateful operator which, for example, maintain accumulated momentum and accumulated squared gradient. However, this way might implicitly introduce assignment semantic into ONNX and therefore breaks SSA assumption.

Below is a script I used to verify the Adagrad function added in this PR's test. It shows that consistent results in TF, Pytorch, and the Adagrad function used to create ONNX tests in this PR.

import numpy as np
import torch
import tensorflow as tf
from torch import optim
from torch import nn
import torch.random as rnd

lr=0.2
lr_decay=0.0
weight_decay=0.0
n = 10
l = 5
X_ = torch.randn(l, n)
Y_ = torch.randn(l, 1)

def apply_adagrad(r, t, x, g, h, norm_coefficient, epsilon, decay_factor):  # type: ignore
    # Compute adjusted learning-rate.
    r_ = r / (1 + t * decay_factor)
    # Add gradient of regularization term.
    g_regularized = norm_coefficient * x + g
    # Update squared accumulated gradient.
    h_new = h + g_regularized * g_regularized
    # Compute ADAGRAD's gradient scaling factors
    h_sqrt = np.sqrt(h_new) + epsilon
    # Apply ADAGRAD update rule.
    x_new = x - r_ * g_regularized / h_sqrt
    return (x_new, h_new)

def show_pytorch():
    X = X_.clone()
    Y = Y_.clone()

    rnd.manual_seed(0)
    model = nn.Sequential(
        nn.Linear(n, 1, bias=False)
    )

    loss_fn = nn.MSELoss(reduction='sum')

    solver = optim.Adagrad(model.parameters(), lr=lr, lr_decay=lr_decay, weight_decay=weight_decay)

    for t in range(8):
        Y_pred = model(X)
        loss = loss_fn(Y_pred, Y)
        model.zero_grad()
        loss.backward()
        solver.step()
        print(loss)

def show_tensorflow():
    rnd.manual_seed(0)
    layer = nn.Linear(n, 1, bias=False)

    X = tf.placeholder('float', shape=[l, n])
    W = tf.Variable(torch.Tensor.numpy(layer.weight.detach()))
    Y = tf.placeholder('float', shape=[l, 1])
    Y_pred = tf.matmul(X, W, transpose_b=True)
    loss = tf.reduce_sum(tf.square(Y - Y_pred))
    optimizer = tf.train.AdagradOptimizer(learning_rate=lr, initial_accumulator_value=0.0000000001)
    minimizer = optimizer.minimize(loss)

    sess = tf.Session()

    init = tf.global_variables_initializer()

    sess.run(init)

    for t in range(8):
        print(sess.run(loss, {X: X_, Y: Y_}))
        result = sess.run([minimizer], {X: X_, Y: Y_})


def show_onnx():
    X = X_.clone()
    Y = Y_.clone()

    rnd.manual_seed(0)
    model = nn.Sequential(
        nn.Linear(n, 1, bias=False)
    )

    loss_fn = nn.MSELoss(reduction='sum')

    solver = optim.Adagrad(model.parameters(), lr=lr, lr_decay=lr_decay, weight_decay=weight_decay)

    for t in range(8):
        Y_pred = model(X)
        loss = loss_fn(Y_pred, Y)
        model.zero_grad()
        loss.backward()

        with torch.no_grad():
            for param in model.parameters():
                new_tensor, new_state = apply_adagrad(r=lr, t=t, x=param.data, g=param.grad.data, h=solver.state[param]['sum'].data,
                                                      norm_coefficient=weight_decay, epsilon=1e-10, decay_factor=lr_decay)
                solver.state[param]['sum'].data = new_state.data
                param.data = new_tensor.data

        print(loss)

show_tensorflow()
print('---------')
show_pytorch()
print('---------')
show_onnx()

[SherlockNoMad]

I echo that having a stateless optimizer is a good design with existing ONNX without breaking SSA assumption.

Different backends can come up with their own implementation on how to handle the new state and update the weights.

[gramalingam]

Hi, Is there a plan to add the complete Adagrad optimization (loop) as another operator? The proposed operator seems only one part of the solution. It would help to understand how the whole thing would work.

[wschin]

Hi, Is there a plan to add the complete Adagrad optimization (loop) as another operator? The proposed operator seems only one part of the solution. It would help to understand how the whole thing would work.

There are two phases of supporting training in ONNX.

1. Training the model with one iteration.
2. Training the model with the whole data set.

Those operators are defined for Phase 1 and will be used in Phase 2 to compose multi-iteration algorithms. Several difficults that we can't go to Phase 2 directly are

1. Defining a full training algorithm is similar to writing a program.
2. Defining a full training algorithm may put extra constraints on the inference stage. For example, if a written training algorithm wants to run 100 iterations, user will have to feed 100 batches into ONNXRuntime.
3. If we can define one training iteration properly, conceptually users can execute multiple iterations.

[postrational]

Overall the proposal looks very good. It seems that this operator could be composed as a Function, could we add a FunctionBody to its definition?

[postrational]

Transforming this to a function may be problematic at this stage due to limitations of the Split and Function.
We should use this op as a test case for expanding the Function and move this op to a function in the future.

[gramalingam]

Many thanks!