Skip to content

Instantly share code, notes, and snippets.

@mratsim
Last active April 28, 2018 08:54
Show Gist options
  • Save mratsim/9d176b7b366fae9b990ff82932756ead to your computer and use it in GitHub Desktop.
Save mratsim/9d176b7b366fae9b990ff82932756ead to your computer and use it in GitHub Desktop.
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
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
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment