Introduce new layer initialization APIs with automatic shape computation

[shadaj]

to be addressed:

• AutoBatchNorm is tied to 2 dimensions + channel inputs, we should handle inputs of any dimension (an HList equivalent (Automatic Requirement Satisfaction?) could help here)
• throw an error if the same key is used twice in the model and offer some alternative for attaching keys to layers in reused chunks

Model Porting Progress:

• LeNet (diff)
• VGG (diff)
• ResNet (diff)
• WordSeg
• BERT
• GPT-2

[saeta]

Couple questions:

1. Why do you separate the definition of the model from the buildModel call that provides the input shape? Could you describe the advantages of this API design vs taking the input right at the start?

2. Why do you use a .then(...) construction, instead of just chaining the calls directly? (Some advantages I can think of for just chaining the calls directly include: (a) IDE autocomplete can work a bit nicer including automatically only having available what layers work on a given input shape, (b) avoids extra syntactical noise. Disadvantages include: when defining a layer, you also want to define an extension method on some protocol to make the autocomplete go.)

Great work so far! I'd definitely make sure to try out a NN with skip connections, such as ResNet sooner rather than later, to ensure your design appropriately takes those into account. (And, of course, as you noted previously, a network that uses weight sharing between layers (aka a network that reuses a layer multiple times).)

[saeta]

CC @dan-zheng @dabrahams

[shadaj]

Responding to @saeta's questions:

1. Passing the input at the beginning/end reduces to a decision about what the "root" node of a layer graph should be. If we pass input at the beginning, the root node for any graph would have to be a node that specifies the input shape. This makes it harder to break down complex graphs into independent chunks. For example, consider a situation where we we want to split out the middle two layers in a graph.

let myModelDefinition = InputNode(shape: ...).then(Dense(...)).then(Dense(...)).then(Dense(...)).then(Dense(...))

If we require input shapes to be specified at the beginning, we have to manually compute the input shape for the sub-graph, which introduces more overhead for the developer and also makes it impossible to re-use the subgraph in another model where the input shape is different.

let mySubModel = InputNode(shape: ...).then(Dense(...)).then(Dense(...))
let myModelDefinition = InputNode(shape: ...).then(Dense(...)).then(mySubModel).then(Dense(...))

If we have users specify the input shape at the end, however, this is solved more naturally because the entire model is composed without being fixed to a specific input shape. We can plug in any input shape at the end and produce a model instance with the appropriate propagated shapes

let mySubModel = Dense(...).then(Dense(...))
let myModelDefinition = Dense(...).then(mySubModel).then(Dense(...))

// elsewhere

var myModelInstance = myModelDefinition.buildModel(inputShape: ...)

2. Couple reasons why I picked .then

• guides type inference to automatically pick the correct Scalar type by looking at the previous layer in the graph (in the example code, we only have to specify Float for the first layer)
• as you mentioned, guides IDE completion
• supports any length of sequential layers since .then can build up a linked-list of types (compared to Sequential which AFAIK is limited in the number of layers since it is just code generated)

Not sure I understood the extension method issue, could you explain further?

[shadaj]

Also, VGG16 works now, but before I proceed to implement more models with this API I'd like to discuss the type complexity being introduced by the AutoLayer protocol.

In the VGG16 implementation, for example, when breaking out the model into a function we have to explicitly state the return type. This type contains the full model structure, which results in some very long types. However, it seems that we need this type information in order to return the appropriate InstanceType when buildModel is called. Any ideas on how we could simplify/eliminate these long types?

[shadaj]

Some early benchmark results:
ResNetCIFAR10 (training) - old: 543.36 examples/sec, new: 435.76 examples/sec
ResNetCIFAR10 (inference) - old: 5286.49 examples/sec, new: 5527.64 examples/sec

With x10:
ResNetCIFAR10 (training) - old: 718.85 examples/sec, new: 753.17 examples/sec
ResNetCIFAR10 (inference) - old: 31586.44 examples/sec, new: 31124.57 examples/sec

[8bitmp3]

@shadaj Awesome work—just saw this. Thanks for the efforts. Just wanted to give my 2 cents - I know some of the stuff is marked as TODO ... oh no 🙂 e.g. ...<AutoSequenced<AutoSequenced<AutoSequenced.... in VGG16. So, hopefully the new stuff will be human friendly and readable in the spirit of Swiftlang. 🙌

[saeta]

re:crazy types, I wonder if associated type inference of opaque result types might be able to help us out here.

[shadaj]

The last commit (4a9719a) sets up the initial infrastructure for a LayerModule that uses properties as a way to get automatic type inference for abstracted layer chunks. Right now, explicit types are still needed when there are multiple lines in the layer creation logic, but with swiftlang/swift#32223 we should have full type inference.

[8bitmp3]

Thanks for this @shadaj and other instances of simpler code. var classifier = ... appears friendlier for the folks who use Pythonic frameworks. Also +1 for the separateBuiltAutoLayer struct and I think the naming is in line with other struct - KeyedAutoLayer - so it makes sense.

[shadaj]

Interestingly, running with -cross-module-optimization results in the linker crashing with SIGBUS (Misaligned address error), but this crash doesn't happen on master. Need to investigate further...

[saeta]

I was thinking about your approach and looking over this code again. I think you're right that being able to specify how to build a model and have that be reusable for a variety of different shapes is a good idea. But I feel like there's a lot of boilerplate required to make a new AutoLayer. While we should care very much about the API to compose pre-written layers, I think it'd also be a good idea to think carefully about the API for making new layers as well.

This led me to wondering: is there a way to have less ceremony without losing the separation between model specification and shape specification. I haven't fully thought this through, but I wonder if closures could actually be that.

Said another way, the difference between:

var model = Sequential<Float>(shape: (1, 2, 3))
  .conv(...)
  .flatten
  .dense(output: 10)

and

let modelBuilder = Sequential<Float>()
   .conv(...)
   .flatten
   .dense(output: 10)
var model = modelBuilder.build(shape: (1, 2, 3))

is delayed binding of the shape. But I think we can get the simplicity of the former with the decoupling of the latter by using a closure / function.

func makeModel(shape: Shape) -> Model {
  return ...
}

Maybe said another way: your design of separating the two is essentially curry-ing the model initialization to separate out other hyperparameters from the shape sizes. One commonly used point in design space for this is to return a closure.

Does that make sense?

[saeta]

One potential idea: consider making a dedicated type that is expressible by array literal (https://developer.apple.com/documentation/swift/expressiblebyarrayliteral)

[shadaj]

@saeta hmm, not sure I completely understand how the shape-first approach simplifies layer definitions. Right now, the reason why we have AutoModule instead of just regular functions is for type inference, but functions do work as well. The only issue with functions is that the return type (which is very long) would have to be explicitly included in the code.

I do really like the chained API style to reduce .then calls for simple sequential models. I'll work on getting that added and using that in the examples, which should give us a better idea of the complexity of the end-user experience.

[saeta]

Sorry for not explaining it well. When reading your code, it looked like there were a lot of structs that had an init to store some values, and just had a single other function to execute their behavior. I've highlighted a few of them below. Does that help?

[saeta]

This to me looks just like a closure capture.

[saeta]

This also just looks like a closure capture.

[saeta]

This also looks like a closure capture.

[saeta]

This also appears to be a hand-written closure capture.

[saeta]

This also looks a bit like a closure capture.

[saeta]

This also looks like a closure capture.

(Also below.)

[shadaj]

Ah, that makes a lot more sense now! I think I have an idea that can get the best of both worlds by implementing these layers as functions that return something similar to AutoFunction but with a function that returns a layer instance instead. Will work on implementing this!

[dabrahams]

The lack of any sort of documentation comments makes it very difficult for me to say anything useful review-wise. I can't tell what these components are supposed to be doing.

[saeta]

I recommend we convert this PR to a draft, as we're not ready to actually merge in the approach based on this direction, and instead continue to refine our thinking per tensorflow/swift#515

[shabalind]

Closing for now. Feel free to reopen if there is more progress!