Reduce explicit types required when instantiating the SGD optimizer.

[saeta]

Previously, in order to instantiate the Optimizer, you had to call
type(of: model) and pass that into the Optimizer constructor in order to get
type inference to pick the right type for Model. This could be a little
confusing for new users.

This commit proposes an alternate way to write this:

let optimizer = SGD(learningRate: 0.01, for: model)

The above formulation is clear and readable. It avoids any unnecessary typing
of types, and should be generalizable to the other optimizers. (This is left as
an exercise for the reader. :-D)

We rely on the Swift optimizer to do the right thing such that only a reference
to the model is passed to the SGD.init call so that we don't pay for the cost
of a full model copy (which could eventually be very expensive). This is deemed
to be a reasonably safe assumption, especially given the ABI standardization in
Swift 5.

[lattner]

Awesome, so much nicer and less foreign, thanks!

Random question for @rxwei and @dan-zheng - would marking the argument explicitly as __shared avoid the possibility of a copy?

[rxwei]

would marking the argument explicitly as __shared avoid the possibility of a copy?

Yep, that is exactly what I was going to suggest!

[rxwei]

Thanks! This looks like a pretty idiomatic use. My other suggestion is to apply this to all other optimizers at once -- they are still using metatypes right now. It would be quite a hassle to update broken models (and toolchains) in different patches.

[jekbradbury]

Any reason not to put the model argument first? I’m OK either way, but having it first would have been my initial intuition.

[rxwei]

Any reason not to put the model argument first? I’m OK either way, but having it first would have been my initial intuition.

I think that intuition is definitely plausible, but here are my reasons not to make for: model the first argument.

1. Per Swift API Design Guidelines, the first argument to an initializer call should not form a phrase starting with the base name. Therefore SGD(for: model, learningRate: ...) would be somewhat not idiomatic. If we made it look like SGD(model: model, learningRate: ...), then the argument label model: implies to me that SGD would hold a copy of model, which is not the case. We want to make it look like SGD is constructed for a model type, instead of owning a model.
2. Making for: model the last argument also helps reduce the salience and the necessity of passing the instance of a model and convey that it's only there for type inference. We still want an initializer without needing this model instance argument, but that can be added in a separate patch.

[rxwei]

There's actually a much more serious issue: Have you actually tested it to make sure this really works?

I don't think this would train at all. Floating point literal 0.02 defaults to Double, so the Scalar associated type in SGD would be double. Then it'll try to get key paths to Tensor<Double>, which do not exist at all in a model where all parameters are Tensor<Float>. So my hypothesis is that this will compile but definitely fail to train. This also exposes how confusing it would be to not make optimizers take a currency scalar type explicitly or not make the model provide one.

Potential solutions:

1. Make the Layer protocol actually require a Scalar associated type as the currency type, but that would be really unusable because it would trigger a lot of protocol conformance errors when the user didn't specify a Scalar type alias.
2. Still require a scalar type explicitly in the optimizer. The way to do that is to require a "scalarType" metatype argument as before: When it's specified as an argument, everything generic parameter gets inferred; when it's not specified as an argument, the compiler would expect it to be provided as a generic type parameter. The crucial thing here is to make Scalar not be inferred from the learningRate: argument -- I thought about this earlier, and this is why I added an explicit scalarType: argument in the first place.

My conclusion is that our optimizers (without undergoing a massive change in all building blocks) really need to be explicit about what the parameter type is that they are optimizing. Given all the trade-offs, making the scalar type explicit on the optimizer side seems to be the most flexible approach that does not come at the cost of making mixed precision models hard to express.

[dan-zheng]

I agree with solution 2: with the current design, scalar type should be made explicit for optimizers.

[saeta]

Yup, you're right that what happens is that it doesn't train. (I didn't notice because it compiled, and there was a previously failing test.)

I've filed TF-302 to follow up with the larger usability issue, and fixed up the API and tests so they're all passing now.

[saeta]

Whoops, just saw @rxwei 's response to @jekbradbury 's comment about ordering (after making the change to the code). Sorry! I'd like to understand this a bit further. Maybe we can all discuss this later today?

[rxwei]

The argument ordering is not a big deal as it's mainly about style -- sure, we can chat about it today. But I think it's critical to get the parameter type issue resolved before discussing other minor issues.

[rxwei]

This is breaking a lot of existing code that use Optimizer<Model, Float>(...). I'd suggested sending another PR that adds the original initializers back, treat those as the canonical initializers, and make the init(for:...) version call those canonical ones.

[marcrasi]

This is breaking a lot of existing code that use Optimizer<Model, Float>(...). I'd suggested sending another PR that adds the original initializers back, treat those as the canonical initializers, and make the init(for:...) version call those canonical ones.

I'll make a PR that does this!