Create a neural net Domain Specific Language

[mratsim]

Currently it is very cumbersome to model a neural network as we have to initialized weight and bias manually. For the MNIST example we want to pass from that:

Old API

# Config (API is not finished)
let
  # We randomly initialize all weights and bias between [-0.5, 0.5]
  # In the future requires_grad will be automatically set for neural network layers

  cv1_w = ctx.variable(
    randomTensor(20, 1, 5, 5, 1'f32) .- 0.5'f32,    # Weight of 1st convolution
    requires_grad = true
    )
  cv1_b = ctx.variable(
    randomTensor(20, 1, 1, 1'f32) .- 0.5'f32,       # Bias of 1st convolution
    requires_grad = true
    )

  cv2_w = ctx.variable(
    randomTensor(50, 20, 5, 5, 1'f32) .- 0.5'f32,   # Weight of 2nd convolution
    requires_grad = true
    )

  cv2_b = ctx.variable(
    randomTensor(50, 1, 1, 1'f32) .- 0.5'f32,       # Bias of 2nd convolution
    requires_grad = true
    )

  fc3 = ctx.variable(
    randomTensor(500, 800, 1'f32) .- 0.5'f32,       # Fully connected: 800 in, 500 ou
    requires_grad = true
    )

  classifier = ctx.variable(
    randomTensor(10, 500, 1'f32) .- 0.5'f32,        # Fully connected: 500 in, 10 classes out
    requires_grad = true
    )

proc model[TT](x: Variable[TT]): Variable[TT] =
  # The formula of the output size of convolutions and maxpools is:
  #   H_out = (H_in + (2*padding.height) - kernel.height) / stride.height + 1
  #   W_out = (W_in + (2*padding.width) - kernel.width) / stride.width + 1

  let cv1 = x.conv2d(cv1_w, cv1_b).relu()      # Conv1: [N, 1, 28, 28] --> [N, 20, 24, 24]     (kernel: 5, padding: 0, strides: 1)
  let mp1 = cv1.maxpool2D((2,2), (0,0), (2,2)) # Maxpool1: [N, 20, 24, 24] --> [N, 20, 12, 12] (kernel: 2, padding: 0, strides: 2)
  let cv2 = mp1.conv2d(cv2_w, cv2_b).relu()    # Conv2: [N, 20, 12, 12] --> [N, 50, 8, 8]      (kernel: 5, padding: 0, strides: 1)
  let mp2 = cv2.maxpool2D((2,2), (0,0), (2,2)) # Maxpool1: [N, 50, 8, 8] --> [N, 50, 4, 4]     (kernel: 2, padding: 0, strides: 2)

  let f = mp2.flatten                          # [N, 50, 4, 4] -> [N, 800]
  let hidden = f.linear(fc3).relu              # [N, 800]      -> [N, 500]

  result = hidden.linear(classifier)           # [N, 500]      -> [N, 10]

# Stochastic Gradient Descent (API will change)
let optim = newSGD[float32](
  cv1_w, cv1_b, cv2_w, cv2_b, fc3, classifier, 0.01f # 0.01 is the learning rate
)

to something like that:

Potential new API

# Config
# Assuming MNIST input [batch_size, color, height, width]
# For N-sized batches: [N, 1, 28, 28]

network model_name:
  context: ctx
  input_shape: [1, 28, 28]              # Real shape [N, 1, 28, 28]
  layers:
    cv1: Conv2D(20, 5, 5)               # Output shape [N, 20, 24, 24] (kernel 5x5, padding 0, stride 1)
    mp1: MaxPool2D((2,2), (0,0), (2,2)) # Output shape [N, 20, 12, 12] (kernel 2X2, padding 0, stride 2)
    cv2: Conv2D(50, 5, 5)               # Output shape [N, 50,  8,  8] (kernel 5x5, padding 0, stride 1)
    mp2: MaxPool2D((2,2), (0,0), (2,2)) # Output shape [N, 50,  4,  4] (kernel 2X2, padding 0, stride 2)
                                        # Flatten      [N, 800]
    hidden: Linear(500)                 # Output shape [N, 500]
    classifier: Linear(10)              # Output shape [N, 10]
  forward x:
    x.cv1.relu.mp1.cv2.relu.mp2.flatten.hidden.relu.classifier

# Stochastic Gradient Descent
let optim = newSGD[float32](model_name, 0.01f)

Importantly, similar to Keras only the output shape is defined within the layers: Linear(500) instead of [500, 800], this greatly ease the usage as it is inferred depending on the sequence of layers.

Implementation

Concretely that means that the network macro should produce something similar to the following:

type ModelName = object
  ctx: Context
  cv1_weight: Variable[subtype(ctx)]
  cv1_bias: Variable[subtype(ctx)]
  mp1: tuple[kernel, padding, stride: Size2D]
  cv2_weight: Variable[subtype(ctx)]
  cv2_bias: Variable[subtype(ctx)]
  mp2: tuple[kernel, padding, stride: Size2D]
  hidden_weight: Variable[subtype(ctx)]
  classifier_weight: Variable[subtype(ctx)]

proc newModelName(): ModelName =
  result.cv1_weight = ctx.variable(
    randomTensor(20, InputSize_1, 5, 5, 1'f32) .- 0.5'f32,
    requires_grad = true
    )
  result.cv1_bias = ctx.variable(
    randomTensor(20, 1, 1, 1'f32) .- 0.5'f32,
    requires_grad = true
    )

  result.cv2_weight = ctx.variable(
    randomTensor(50, InputSize_2, 5, 5, 1'f32) .- 0.5'f32,
    requires_grad = true
    )

  result.cv2_bias = ctx.variable(
    randomTensor(50, 1, 1, 1'f32) .- 0.5'f32,
    requires_grad = true
    )

  ...
  # Note that we must automatically determine the input_size1 and input_size2
  # Depending on "input_shape" and the previous layers

proc forward(self: ModelName, x: Variable[T]): Variable[T] =
  template cv1(x: Variable[T]): Variable[T]:
    x.conv2d(self.cv1_weight, self.cv1_bias)
  template cv2(x: Variable[T]): Variable[T]:
    x.conv2d(self.cv2_weight, self.cv2_bias)

  ...

  x.cv1.relu.mp1.cv2.relu.mp2.flatten.hidden.relu.classifier

Things to keep in mind:

• The model must be serializable so we can save a trained model to disk
• It should be possible to mix stateless functions like relu, sigmoid and flatten in the forward proc.
• The forward proc should allow a variadic number of inputs. For example to combine video and audio
• We can have temporary values in the forward proc

References:

• Neural-Network in Nim syntax: https://github.com/t8m8/Neural-Network-in-Nim/blob/4bd06420fd7c57afb665c6b4e5148527eda6a410/examples/xor.nim#L14-L21 (includes input and output shape)
• Keras syntax (Sequential and Functional): https://keras.io/getting-started/functional-api-guide/ (only output shape, input is inferred)
• PyTorch syntax: https://github.com/pytorch/examples/blob/dcdabc22b305d2f2989c6f03570dfcd3919e8a5b/mnist/main.py#L52-L68 (input and output shape)
• Tensorflow eager syntax: https://www.tensorflow.org/get_started/eager (infers input shape when going through Keras)
• Grenade (Haskell): https://github.com/HuwCampbell/grenade

type MNIST
  = Network
    '[ Convolution 1 10 5 5 1 1, Pooling 2 2 2 2, Relu
     , Convolution 10 16 5 5 1 1, Pooling 2 2 2 2, Reshape, Relu
     , FullyConnected 256 80, Logit, FullyConnected 80 10, Logit]
    '[ 'D2 28 28, 'D3 24 24 10, 'D3 12 12 10, 'D3 12 12 10
     , 'D3 8 8 16, 'D3 4 4 16, 'D1 256, 'D1 256
     , 'D1 80, 'D1 80, 'D1 10, 'D1 10]

randomMnist :: MonadRandom m => m MNIST
randomMnist = randomNetwork

Future

The training loop and validation can be also syntactic sugar-ed with fit and predict

[mxbi]

I am a big fan of the potential new API syntax in your post. Coming from keras, a 'wishlist' of things that would be extremely useful here for a NN API:

• Something analogous to model.summary(), which shows the shapes/weight count of each layer and various other properties. The ability to plot a flowchart of the model would give you a one-up over other libraries.
• Being able to set properties on individual layers such as freezing weights (in Keras) or LR multipliers (not implemented in Keras).
• Sharing weights between layers or applying the same layers to two inputs then merging (eg. for siamese nets)
• Being able to define new models as supersets/subsets of existing models (useful for pretrained models).

[mratsim]

1. Model summary + flowchart

Yes for model.summary, is a flowchart this?

Taken from my multi-attention neural network here

def Net(num_feats, seq_len, num_hidden, num_outputs):
    x = Input(shape=(seq_len, num_feats))

    # Encoder RNNs
    enc = CuDNNGRU(seq_len,
                   return_sequences=True,
                   stateful = False,
                   name = 'Encoder_RNN')(x)
    
    # Attention decoders (lag features)
    attention_0d = attention_n_days_ago(enc, 0)
    attention_1d = attention_n_days_ago(enc, 1)
    attention_2d = attention_n_days_ago(enc, 2)
    attention_4d = attention_n_days_ago(enc, 4)
    attention_1w = attention_n_days_ago(enc, 7)
    attention_2w = attention_n_days_ago(enc, 14)
    attention_1m = attention_n_days_ago(enc, 30)
    attention_2m = attention_n_days_ago(enc, 60)
    attention_1q = attention_n_days_ago(enc, 92)
    attention_6m = attention_n_days_ago(enc, 184)
    attention_3q = attention_n_days_ago(enc, 276)
    attention_1y = attention_n_days_ago(enc, 365)
    
    att = Concatenate(name='attns_cat', axis = 1)([attention_0d,
                                                   attention_1d,
                                                   attention_2d,
                                                   attention_4d,
                                                   attention_1w,
                                                   attention_2w,
                                                   attention_1m,
                                                   attention_2m,
                                                   attention_1q,
                                                   attention_6m,
                                                   attention_3q,
                                                   attention_1y])
    
    # How to merge? concat, mul, add, use Dense Layer or convolution ?
        
    att = Dense(seq_len, activation=None, name='Dense_merge_attns')(att)
    # att = Lambda(lambda x: softmax(x, axis = 1),
    #              name='Dense_merge_3D_softmax')(att) # Flatten along the concat axis
    
    # Decoder RNN
    dec = CuDNNGRU(num_hidden,
                   return_sequences=False,
                   stateful = False,
                   name='Decoder_RNN')(att)

    # Regressor
    # Note that Dense is automatically TimeDistributed in Keras 2
    out = Dense(num_outputs, activation=None,
                name = 'Classifier')(dec) # no activation for regression
    
    model = Model(inputs=x, outputs=out)
                          
    model.compile(loss= root_mean_squared_error, optimizer = optim)
    return model

2. Individual updates (freezing/ LR multiplier)

Freezing weights is planned. Individual LR multiplier would be easy.

  VariableObj[TT] = object
    context*: ContextPtr[TT] # This is the computational graph
    value*: TT                          # This is the value Tensor
    grad*: TT                           # This is the gradient Tensor
    requires_grad*: bool         # This tells if gradient is needed
    ## new fields
    frozen: bool
    lr_multiplier: float32 or float64

3. Weights sharing

Obviously, this is my main grip against Keras and why I like PyTorch, especially when you need RNNs. Weights are shared in this case:

# Assuming MNIST double inputs [batch_size, color, height, width]
# For N-sized batches: [N, 1, 28, 28], [N, 1, 28, 28]

network model_name:
  context: ctx
  input_shape: [1, 28, 28], [1, 28, 28] # Real shape [N, 1, 28, 28], [N, 1, 28, 28]
  layers:
    cv1: Conv2D(20, 5, 5)               # Output shape [N, 20, 24, 24] (kernel 5x5, padding 0, stride 1)
    mp1: MaxPool2D((2,2), (0,0), (2,2)) # Output shape [N, 20, 12, 12] (kernel 2X2, padding 0, stride 2)
    hidden: Linear(500)                 # Output shape [N, 500]
    classifier: Linear(10)              # Output shape [N, 10]
  forward x, y:
    let foo1 = x.cv1.relu.mp1
    let foo2 = y.cv1.relu.mp1
    let merge = merge(foo1, foo2)
    result = merge.flatten.hidden.relu.classifier

Note: I have to figure out how to share state and weights for RNNs especially for seq2seq.

4. Models as subsets/supersets

I also need and want that. I still don't know how in practice but since all networks will have a standard forward proc I will probably use that. In any case I don't think you can define ResNets properly without this.

[mxbi]

1. Yep, this is the flowchart I mean.

3. This sort of weight sharing can be done in keras now in the same way as your example here (slightly-pseudocode):

cv1 = Conv2D(20, 5, 5)
mp1 = MaxPool2D((2, 2), (0, ), (2, 2))
relu = Activation('relu')

foo1 = mp1(relu(cv1(x)))
foo2 = mp1(relu(cv1(y))) # Shared weights

I think the way you have it implemented now is good.

[mratsim]

I've been hacking at this and I got a parser working for the following syntax.

network FooNet:
  context: ctx
  layers:
    x:   Input([1, 28, 28])
    y:   Input([1, 28, 28])
    cv1: Conv2D(x.oshape, 20, 5, 5)
    mp1: MaxPool2D(cv1.oshape, (2,2), (0,0), (2,2))
    classifier: Linear(mp1.oshape.flatten, 10)
  forward x, y:
    x.cv1.relu.mp1.flatten.classifier

Any number of inputs/outputs are supported. to do weight sharing you only need to reuse the same name in the forward part (assuming a merge exist):

network FooNet:
  context: ctx
  layers:
    x:   Input([1, 28, 28])
    y:   Input([1, 28, 28])
    cv1: Conv2D(x.oshape, 20, 5, 5)
    mp1: MaxPool2D(cv1.oshape, (2,2), (0,0), (2,2))
    classifier: Linear(mp1.oshape.flatten, 10)
  forward x, y:
    let z1 = x.cv1.relu.mp1.flatten.classifier
    let z2 = y.cv1.relu.mp1.flatten.classifier
    return merge(z1, z2, Sum)

The forward section is completely free like in dynamic frameworks and supports arbitrary Nim syntax, including temporary variables.

The hard part was Keras automatic input shape inference from the sequence of layers. I checked in PyTorch and Chainer code, and also Keras code from 2015, I know how to do it for a static graph but I don't see how to do it for a dynamic graph that can allow arbitrary code for example if there is an arbitrary kaggle_flatten function like:

network FooNet:
  context: ctx
  layers:
    x:   Input([1, 28, 28])
    y:   Input([1, 28, 28])
    cv1: Conv2D(x.oshape, 20, 5, 5)
    mp1: MaxPool2D(cv1.oshape, (2,2), (0,0), (2,2))
    classifier: Linear(???????, 10)
  forward x, y:
    let z1 = x.cv1.relu.mp1.kaggle_flatten.classifier
    let z2 = y.cv1.relu.mp1.kaggle_flatten.classifier
    return merge(z1, z2, `*`)

Instead each layers keep track of its input, compute its output shape to avoid manual computation like it's done in PyTorch.

key: y, input: [1, 28, 28], output: [1, 28, 28]
key: cv1, input: [1, 28, 28], output: [20, 24, 24]
key: x, input: [1, 28, 28], output: [1, 28, 28]
key: mp1, input: [20, 24, 24], output: [20, 12, 12]
key: classifier, input: [2880], output: 10

[manguluka]

@mratsim Many thanks for your work on Arraymancer.

Is the WIP parser publicly available?

[mratsim]

@manguluka here you go: https://gist.github.com/mratsim/9d176b7b366fae9b990ff82932756ead#file-wip_nnet_dsl-nim.

You need the src/tensor/backend/metadataArray, for now it only validates and parses the configuration but does not create the model.

[mratsim]

And the main syntax is done 🔥

Here is a MNIST with 2 Conv layers + 1 hidden linear layer.

import ../src/arraymancer, random

randomize(42)

let
  ctx = newContext Tensor[float32] # Autograd/neural network graph
  n = 32                           # Batch size

let
  x_train = read_mnist_images("build/train-images.idx3-ubyte").astype(float32) / 255'f32
  X_train = ctx.variable x_train.unsqueeze(1)

  y_train = read_mnist_labels("build/train-labels.idx1-ubyte").astype(int)

  x_test = read_mnist_images("build/t10k-images.idx3-ubyte").astype(float32) / 255'f32
  X_test = ctx.variable x_test.unsqueeze(1)
  y_test = read_mnist_labels("build/t10k-labels.idx1-ubyte").astype(int)

network ctx, DemoNet:
  layers:
    x:          Input([1, 28, 28])
    cv1:        Conv2D(x.out_shape, 20, 5, 5)
    mp1:        MaxPool2D(cv1.out_shape, (2,2), (0,0), (2,2))
    cv2:        Conv2D(mp1.out_shape, 50, 5, 5)
    mp2:        MaxPool2D(cv2.out_shape, (2,2), (0,0), (2,2))
    hidden:     Linear(mp2.out_shape.flatten, 500)
    classifier: Linear(500, 10)
  forward x:
    x.cv1.relu.mp1.cv2.relu.mp2.flatten.hidden.classifier

let model = ctx.init(DemoNet)
let optim = model.optimizerSGD(learning_rate = 0.01'f32)

# Learning loop
for epoch in 0 ..< 5:
  for batch_id in 0 ..< X_train.value.shape[0] div n: # some at the end may be missing, oh well ...
    # minibatch offset in the Tensor
    let offset = batch_id * n
    let x = X_train[offset ..< offset + n, _]
    let target = y_train[offset ..< offset + n]

    let clf = model.forward(x)
    let loss = clf.sparse_softmax_cross_entropy(target)

    if batch_id mod 200 == 0:
      # Print status every 200 batches
      echo "Epoch is: " & $epoch
      echo "Batch id: " & $batch_id
      echo "Loss is:  " & $loss.value.data[0]

    loss.backprop()
    optim.update()

  # Validation (checking the accuracy/generalization of our model on unseen data)
  ctx.no_grad_mode:
    echo "\nEpoch #" & $epoch & " done. Testing accuracy"

    # Validation by batches of 1000 images
    var score = 0.0
    var loss = 0.0
    for i in 0 ..< 10:
      let y_pred = model.forward(X_test[i*1000 ..< (i+1)*1000, _]).value.softmax.argmax(axis = 1).indices.squeeze
      score += accuracy_score(y_test[i*1000 ..< (i+1)*1000], y_pred)

      loss += model.forward(X_test[i*1000 ..< (i+1)*1000, _]).sparse_softmax_cross_entropy(y_test[i*1000 ..< (i+1)*1000]).value.data[0]
    score /= 10
    loss /= 10
    echo "Accuracy: " & $(score * 100) & "%"
    echo "Loss:     " & $loss
    echo "\n"

Sharing weights should work. Missing:

• - Freezing layers. Actually you can do:
  • model.cv1.weights.requires_grad = false; model.cv1.bias.requires_grad = false but that could be better
• - individual learning rate
• - Initialization (Xavier Glorot, Kaiming He)
• - model summary

[mratsim]

Following a request on Reddit, everything that can change the shape/topology will be available as layers, notably Flatten and Reshape. That will make it easier to print a graph too I think.

I don't think it's useful to put activations (relu/tanh/sigmoid/softmax) as part of the layers though, except for RNN internal. I'll see when i add the RNNs.