Framework: Move Ops parameter to call method where possible

[rnett]

I'd like to move the Ops parameters of framework classes to the call method, where possible. This is primarily for Kotlin interop, but has a few other benefits as well. It won't be possible for stateful classes (Metrics, Optimizers), but should be possible for most, as far as I can tell (Initializers, Activations, Losses). I'm going to use Losses as a standin for all 3 for my examples.

1. Kotlin interop w/ @FunctionInterface. If we do this, we can then define new losses like Loss{ tf, x -> tf.stuff(x) } or pass lambdas to methods that take losses. This is very nice for layers, where we might have something Keras-like Layer(activation=ReLU()) but want to replace it with something custom.
2. Re-use of objects. Currently, losses create any subsequent calls in the same scope as their first call. That means if it's called inside a sub-scope, including device ones, it ignores the scope. This is somewhat expected, but not ideal. It also causes further issues if we use Ops for (eager) tensor lifetime management, which has been suggested and is something I'd like to do (it's easy enough to make a long-lived copy of the initial scope, but then the tensors created in the call methods live forever, and the framework classes need to be closable).
3. Passing configs, like Keras. In Keras, if you have a activation or loss that requires some parameters, you can pass it to a layer like Layer(activation=LeakyReLU(alpha=0.3)). I expect this will be common with our API, as well. Currently, this runs into the above issue w/ scoping, and prevents you from passing activations (or losses) from scopes that don't have an Ops instance available.

I'd look at having the stateful classes take Ops in call as well, and only using the constructor ops for initializing state. This works better with scoping and lifetimes as mentioned above.

[karllessard]

@JimClarke5 , we've been discussing Ryan and I about this and would like to revive this proposal, whenever you're ready

[JimClarke5]

Actually, this may help with some of the stuff I am doing for Model. The Model creates the graph, so it is difficult to create an Optimizer first, until the Model does this. Right now I have Model using Lambda functions to create Loss, Metrics and Optimizers after the Model creates its Graph/Ops.

I am open to a way to defer setting the Ops (or whatever) later on in the Object's life cycle.

BTW: I think we should revisit Optimizers to be consistent with the rest of the framework packages.

[rnett]

Optimizers will probably want to (only?) support the new variables, so you may want to wait for that. It would be nice if there was an option to pass in a list of variables to optimize rather than checking the graph, for eventual eager support.

[JimClarke5]

Some of the classes create internal Variable's and such before the call method. I would vote to use setTF(Ops tf) or init(Ops tf) to initialize the classes and take Ops tf out of the ctors. We could optionally pass tf to call methods, which would call setTF if tf is not already set. One issue is, if the class creates Variable's, then a session.run() needs to be run to initialize the variable before it is used the first time. We might be able to use control dependencies to do this, but the control dependency can only exist when the Variable is first used, and not after.

[rnett]

I agree w/ init, I was doing something similar in Kotlin w/ delegates, so supporting that is nice. It might be nice to have a overloadable, class-scoped init but also a onInit that takes a lambda to call on initialization for things like initing variables.

Something like:

Variable<TFloat32> w = newVariable("w", Shape.of(10, 10));
onInit(() -> w.initWith(...));

Most of the variable stuff should be handled by the new API (using init scopes), once the gradient makes it into tensorflow core, but if you want to apply a initializer to a variable or something it's still good to have that step.

[rnett]

This applies to Layer, as well.

[JimClarke5]

Should we create an interface, Initable ?

public interface Initable {
    public void init(Ops tf);
    public void onInit(Consumer<Ops> onInit);
}

[JimClarke5]

Also, I think we would have to do all the ctors in framework that take an Ops.

[rnett]

It would be nice to not have to expose it to the user, i.e. something like:

abstract class Activation {
  private List<Consumer<Ops>> initers = new ArrayList<>();
  private boolean inited = false;
  protected final void onInit(Consumer<Ops> onInit){
    initers.add(opInit);
  }

  protected void init(Ops tf){}

  public final Operand<TFloat32> call(Ops tf, Operand<TFloat32> x){
    if(!inited){
      init(tf);
      initers.forEach((x) -> x.call(tf));
      inited = true;
    }
    return doCall(tf, x);
  }

  abstract Operand<TFloat32> doCall(Ops tf, Operand<TFloat32> x);
}

where the init stuff gets called on the first call. We could probably put the actual init stuff in an abstract class though.

[JimClarke5]

So with the call proposal, what is the expected behavior the 2nd time call is invoked?

[rnett]

Right, duh, there should be a flag, edited it. The init logic should probably be moved to a synchronized method.

[JimClarke5]

I do not favor passing Ops in the call or similar method, but rather in an init(Ops method) that is invoked one time. This is primarily due to

1. The number of times a call or similar method might be invoked, especially during training.
2. The off chance that a different Ops is passed after another Ops was passed the first time, representing distinct graphs.
3. Handling of initializers that need to be invoked on the initialization pass (e.g. Metric variables).

My preference is Op init(Ops), so that the code can call session.run(instance.init(tf) to do the initialization before anything else.

[rnett]

I don't see the issue with 1, it's still re-creating all the operations, just in the same scope or a new one. And since you're passing in operands you almost always have the Ops available.

2 is exactly why I wanted Ops in the call method, since with Keras, you can do something like create a ReLU(0.3) object and then re-use it in different layers and even models. Same for Layers (which share weights, when used like this). This doesn't work if we always use the same Ops instance in call. The different graphs is a bit complex, too: oftentimes, when using functions, you will be able to say create an initialized variable in the parent graph, and use it in the function. That is supported and rather essential. Using things from eager sessions in graphs is (somewhat) supported, too. Using an incompatible Ops in a call would produce an exception, the same as if you used a Operand from said Ops, and leave doing so up to the ExecutionEnvironment implementation (since some will support it). Additionally, some framework classes work fine regardless of the init Ops since they don't store anything in initialization (i.e. most Activations).

For 3, I envisioned having Ops in call and init, so that you would still be able to do initializers and whatnot (3). You could use Ops.addInit or eventually initScope to run things on Session init, or not (and users of the class wouldn't have to register it themselves).

Also, I would like to separate the "Ops in call" from "init" a little bit. It's much nicer for the Kotlin API if we can make some of the commonly customized Framework classes (Activation, Loss) functional interfaces, and for the most part they don't need initialization, so you could have something like Activation with just the call method and StatefulActivation with the init handling (and perhaps an Initable base class, but that's an implementation detail). That's not possible for things like Optimizers (or Metrics?), but you won't be commonly passing custom ones there anyways.

[JimClarke5]

FYI, It is not just call(), in Metric there is updateStateList, result, and resetStates.

[rnett]

Hmm, yeah, it seems like there's a definite dichotomy between stateful/graph linked classes like Optimizer and Metric and stateless ones like Activation, Loss, Regularizer, etc. I would be fine with just doing the Ops in call for the stateless ones, since you're not going to be able to pass lambdas for the stateful ones anyways, and they aren't customized as much. You could parameterize the state, but that seems like too much trouble to be worth doing.