neural_net_dsl.nim

import macros, ../src/arraymancer

# Transforms

dumpTree:
  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

# Into

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

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

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

    result.cv2_bias = result.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

parserNeuralNet.nim

import breeze, macros

type NetworkParserState = ref object
  ## State of the parser of the network configuration
  context: NimNode
  shape: NimNode
  layers: NimNode
  forward_body: NimNode
  forward_template: NimNode
  asserts: NimNode
  typeDecl: NimNode

proc parseSections(config: NimNode): tuple[context, input_shape, layers, forward: NimNode] =

  template printError =
    error:
      lineInfo(section) &
        ": unknown neural network configuration section \"" &
        $section[0] & "\""

  for section in config:
    if section.kind == nnkCall:
      if eqIdent(section[0], "context"):
        result.context = section[1]
      elif eqIdent(section[0], "input_shape"):
        result.input_shape = section[1]
      elif eqIdent(section[0], "layers"):
        result.layers = section[1]
      else:
        printError()
    elif section.kind == nnkCommand:
      if eqIdent(section[0], "forward"):
        result.layers = section[1]
      else:
        printError()
    else:
        printError()

macro network(model_name: untyped, config: untyped): untyped =

  # 1. Separate the configuration into the context, shape, layers, forward part
  let sections = config.parseSections


network ResNet:
  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
  test: bar

wip_nnet_dsl.nim

import macros, tables, hashes

import ../src/tensor/backend/metadataArray

proc toMetadataArray*(s: MetadataArray): MetadataArray {.inline.} =
  # no-op
  s

proc flatten*(s: MetadataArray): MetadataArray {.inline.}=
  result.len = 1

  result.data[0] = 1

  for val in s:
    result.data[0] *= val

static: echo "\n###################################\n"

proc hash(x: NimNode): Hash =
  assert x.kind == nnkIdent
  result = hash($x)

type
  LayerKind = enum
    lkInput, lkConv2D, lkLinear, lkMaxPool2D

  LayerTopology = object
    ## Describe a layer topology
    ishape, oshape: NimNode # Input and output shape
    case kind: LayerKind
    of lkConv2D:
      c2d_padding, c2d_strides: NimNode
    of lkMaxPool2D:
      m2d_kernel, m2d_padding, m2d_strides: NimNode
    else:
      discard

  Layer = object {.inheritable.}
    ## Holds trainable parameters of each layer
    ## Also holds metadata to be able to print the Network graph

  Conv2DLayer{.final.} = object of Layer
    weight: int
    bias: int

  LinearLayer{.final.} = object of Layer
    weight: int
    bias: int

  #################################################

  TopoTable = Table[NimNode, LayerTopology]

  NetworkSections = ref object
    context, layers, forward: NimNode

  NeuralGraphVM = ref object
    ## Virtual Machine that will go through the network topology

    # Running state
    # We traverse the abstract syntax tree with
    #   - each layer being an operation
    #   - the input variables, intermediate variables
    #     and intermediate layer results being operand
    operands_stack: seq[NimNode]  # Stack of operands for a an expression. Resets after each line
    operators_stack: seq[NimNode] # Stack of operators for a an expression. Resets after each line
    topoTable: TopoTable          # Retrieve the layer properties.
                                  # Maps a NimNode of nnkIdent to the corresponding input and output shape

    # Outputs
    context: NimNode                                              # copy-paste of the context
    modelWeights: seq[tuple[field: NimNode, value: NimNode]]      # must go through the forward, computing weights when needed.
    forward_proc: NimNode                                         # proc forward(x, y)
    forward_body: NimNode                                         # Copy paste of the body "forward x, y"
    forward_templates: seq[tuple[name: NimNode, body: NimNode]]   # We must inster some templates to replace
                                                                  # template cv1(x: Variable[T]): Variable[T] =
                                                                  #   x.conv2d(self.cv1_weight, self.cv1_bias)
    forward_asserts: NimNode

proc oshape_conv2d(ishape, kernel: MetadataArray, padding, strides: tuple[h, w: int]): MetadataArray =
  ## Each dimension of the (nbDims-2)-D images of the output tensor is computed as followed:
  ##   outputDim = 1 + ( inputDim + 2*pad - (((filterDim-1)*upscaleA)+1) )/ convolutionStride;

  ## Input and result are of shape [C, H, W]
  ## Kernel of shape [C_out, C_in, h, w]

  # Reminder we don't consider N (bath size) in the topology
  template kH: int = kernel[2]
  template kW: int = kernel[3]
  template pH: int = padding.h
  template pW: int = padding.w
  template sH: int = strides.h
  template sW: int = strides.w

  template iH: int = ishape[1]
  template iW: int = ishape[2]
  template dH: int = 1 # dilation # TODO
  template dW: int = 1 # dilation

  result.len = 3
  result[0] = kernel[0]                                    # C
  result[1] = 1 + (iH + 2*pH - (((kH-1) * dH) + 1) div sH) # H
  result[2] = 1 + (iW + 2*pW - (((kW-1) * dW) + 1) div sW) # W

proc oshape_maxpool2d(ishape: MetadataArray, kernel, padding, strides: tuple[h, w: int]): MetadataArray =

  # Reminder we don't consider N (bath size) in the topology
  template C: int = ishape[0]
  template H: int = ishape[1]
  template W: int = ishape[2]

  template kH: int = kernel.h
  template kW: int = kernel.w
  template pH: int = padding.h
  template pW: int = padding.w
  template sH: int = strides.h
  template sW: int = strides.w

  result.len = 3
  result[0] = C
  result[1] = (H + (2 * pH) - kH) div sH + 1
  result[2] = (W + (2 * pW) - kW) div sW + 1


proc splitSections(config: NimNode): NetworkSections =
  new result
  template unknown =
    error:
      lineInfo(section) &
        ": unknown neural network configuration section \"" &
        $section[0] & "\""

  for section in config:
    if section.kind == nnkCall:
      if eqIdent(section[0], "context"):
        result.context = section[1]
      elif eqIdent(section[0], "layers"):
        result.layers = section[1]
      else:
        unknown()
    elif section.kind == nnkCommand:
      if eqIdent(section[0], "forward"):
        # For forward we copy everything.
        # We have to deal with forward with multiple inputs like "forward x, y, z:"
        # and we will do that later.
        result.forward = section
      else:
        unknown()
    else:
        unknown()

proc isValidLayerSection(section: NimNode): bool =

  # Expected AST - cv1: Conv2D(20, 5, 5)
  # Call
  #   Ident "cv1"
  #   StmtList
  #     Call
  #       Ident "Conv2D"
  #       IntLit 20
  #       IntLit 5
  #       IntLit 5
  (section.kind == nnkCall) and
    (section[0].kind == nnkIdent) and
    (section[1].kind == nnkStmtlist) and
    (section[1].len == 1) and
    (section[1][0].kind == nnkCall)

template unknown(section: Nimnode) =
  error:
    lineInfo(section) &
      ": unknown neural network configuration section \"" &
      $section[0] & "\""

template incorrect(section: Nimnode) =
  error:
    lineInfo(section) &
      ": incorrect neural network configuration section \"" &
      $section[0] & "\""

proc topoFromInput(self: var TopoTable, ident: NimNode, desc: NimNode) =

  # Initializes the ident --> (input, output) topology table with the input shapes

  # Call
  #   Ident "Input"
  #   Bracket
  #     IntLit 1
  #     IntLit 28
  #     IntLit 28

  if desc.len != 2:
    incorrect(desc) ## Placeholder to specify padding stride in the future

  self.add ident, LayerTopology(kind: lkInput,
                                ishape: desc[1],
                                oshape: desc[1])

proc replaceInputNodes(self: TopoTable, ishape: NimNode): NimNode =
  # Args:
  #   - The topology table
  #   - the input shape
  # Returns:
  #   - An AST input shape with "x.oshape" replaced by the actual x.oshape
  #     taken from the topology table

  template letsGoDeeper =
    var rTree = node.kind.newTree()
    for child in node:
      rTree.add inspect(child)
    return rTree

  proc inspect(node: NimNode): NimNode =
    case node.kind:
    of nnkDotExpr:
      if eqIdent(node[1], "oshape"):
        return self[node[0]].oshape
      else:
        letsGoDeeper()
    of {nnkIdent, nnkSym, nnkEmpty}:
      return node
    else:
      letsGoDeeper()

  result = inspect(ishape)

proc topoFromConv2D(self: var TopoTable, ident: NimNode, desc: NimNode) =

  # Call
  #   Ident "Conv2D"
  #   # Kernel (C_out, kH, kW)
  #   IntLit 20
  #   IntLit 5
  #   IntLit 5
  #   # Kernel strides & padding

  var padding, strides: NimNode

  if desc.len > 5:
    incorrect(desc) ## Placeholder to specify padding stride in the future
  else:
    padding = quote do: (0, 0)
    strides = quote do: (1, 1)

  var ishape = self.replaceInputNodes(desc[1])
  ishape = quote do: `ishape`.toMetadataArray

  let
    c_out = desc[2]
    kH = desc[3]
    kW = desc[4]

  let kernel = quote do:
    # C_out, C_in, kH, kW
    [`c_out`, `ishape`[0], `kH`, `kW`].toMetadataArray

  let oshape = quote do:
    oshape_conv2d(`ishape`, `kernel`, `padding`, `strides`)

  self.add ident, LayerTopology(kind: lkConv2D,
                                ishape: ishape,
                                oshape: oshape,
                                c2d_padding: padding,
                                c2d_strides: strides)

proc topoFromMaxPool2D(self: var TopoTable, ident: NimNode, desc: NimNode) =

  # Call
  #   Ident "MaxPool2D"
  #   Par
  #     IntLit 2
  #     IntLit 2
  #   Par
  #     IntLit 0
  #     IntLit 0
  #   Par
  #     IntLit 2
  #     IntLit 2

  if desc.len != 5:
    incorrect(desc) ## Placeholder to specify padding stride in the future

  var ishape = self.replaceInputNodes(desc[1])
  ishape = quote do: `ishape`.toMetadataArray

  let
    kernel = desc[2]
    padding = desc[3]
    strides = desc[4]

  let oshape = quote do:
    oshape_maxpool2d(`ishape`, `kernel`, `padding`, `strides`)

  self.add ident, LayerTopology(kind: lkMaxPool2D,
                                ishape: ishape,
                                oshape: oshape,
                                m2d_kernel: kernel,
                                m2d_padding: padding,
                                m2d_strides: strides)

proc topoFromLinear(self: var TopoTable, ident: NimNode, desc: NimNode) =

  # Call
  #   Ident "Linear"
  #   IntLit 10

  if desc.len != 3:
    incorrect(desc) ## Placeholder to specify padding stride in the future

  var ishape = self.replaceInputNodes(desc[1])
  ishape = quote do: `ishape`.toMetadataArray

  self.add ident, LayerTopology(kind: lkLinear,
                                ishape: ishape,
                                oshape: desc[2])

proc topoFromLayer(self: var TopoTable, ident: NimNode, desc: NimNode) =

  if eqIdent(desc[0], "Conv2D"):
    self.topoFromConv2D(ident, desc)
  elif eqIdent(desc[0], "MaxPool2D"):
    self.topoFromMaxPool2D(ident, desc)
  elif eqIdent(desc[0], "Linear"):
    self.topoFromLinear(ident, desc)
  elif eqIdent(desc[0], "Input"):
    self.topoFromInput(ident, desc)
  else:
    unknown(desc)

proc topoFromLayers(self: var TopoTable, layers: NimNode) =

  ## Add all layers and their known parameters to the table
  #

  for section in layers:
    if section.isValidLayerSection:
      assert section[0] notin self
      self.topoFromLayer(section[0], section[1][0])
    else:
      unknown(section)

macro network(model_name: untyped, config: untyped): untyped =

  # 0. Separate the configuration into the context, input_shapes, layers, forward part
  let sections = config.splitSections

  # 1. Initilaize the VM to analyse the neural network Graph.
  #    - Get the input shapes
  #    - Get the layers
  let vm = new NeuralGraphVM
  vm.context = sections.context
  vm.topoTable = initTable[NimNode, LayerTopology]()
  vm.topoTable.topoFromLayers(sections.layers)

  result = newStmtList()

  for k, v in pairs(vm.topoTable):
    let i = v.ishape
    let o = v.oshape
    let ident = $k
    result.add quote do:
      echo "key: " & `ident` & ", input:" & $`i` & ", output:" & $`o`

network FooNet:
  context: ctx
  layers:
    x:   Input([1, 28, 28])                         # Real shape [N, 1, 28, 28]
    y:   Input([1, 28, 28])                         # Real shape [N, 1, 28, 28]
    cv1: Conv2D(x.oshape, 20, 5, 5)                 # Output shape [N, 20, 24, 24] (kernel 5x5, padding 0, stride 1)
    mp1: MaxPool2D(cv1.oshape, (2,2), (0,0), (2,2)) # Output shape [N, 20, 12, 12] (kernel 2X2, padding 0, stride 2)
    classifier:
      Linear(mp1.oshape.flatten, 10)                # Output shape [N, 10]
  forward x, y:
    x.cv1.relu.mp1.flatten.classifier