[NN primitives] unify var parameters vs return values

[mratsim]

The library is inconsistent in nn_primitives for when to use return values and when to use var parameters that will be updated in place.

var parameter

Arraymancer/src/nn_primitives/nnp_linear.nim

Lines 20 to 28 in 92aa82e

	proc linear*[T](input, weight: Tensor[T], bias: Tensor[T], output: var Tensor[T]) {.inline.} =
	# Linear (Dense) forward primitive with bias
	# - input tensor shape [batch_size, in_features]
	# - weight tensor shape [out_features, in_features]
	# - bias tensor shape [batch_size, out_features]
	# Output does not need to be initialized to 0 or the proper shape, data will be overwritten
	# Output is: Y = x * W.transpose + b
	output = input * weight.transpose
	output .+= bias

Arraymancer/src/nn_primitives/fallback/conv.nim

Lines 19 to 20 in 92aa82e

	proc im2col[T]( input: Tensor[T], kernel_size: Size2D,
	padding: Size2D = (0,0), stride: Size2D = (1,1), result: var Tensor[T]) =

Arraymancer/src/nn_primitives/nnp_convolution.nim

Lines 65 to 70 in 92aa82e

	proc conv2d_backward*[T](input, weight, bias: Tensor[T],
	padding: Size2D,
	stride: Size2D,
	grad_output: Tensor[T],
	grad_input, grad_weight, grad_bias: var Tensor[T],
	algorithm = Conv2DAlgorithm.Im2ColGEMM) =

return value

Arraymancer/src/nn_primitives/fallback/conv.nim

Lines 57 to 59 in 92aa82e

	proc col2im*[T](input: Tensor[T], channels, height, width: int,
	kernel_size: Size2D,
	padding: Size2D = (0,0), stride: Size2D = (1,1)): Tensor[T] =

Arraymancer/src/nn_primitives/nnp_convolution.nim

Lines 28 to 31 in 92aa82e

	proc conv2d*[T](input, weight, bias: Tensor[T],
	padding: Size2D = (0,0),
	stride: Size2D = (1,1),
	algorithm = Conv2DAlgorithm.Im2ColGEMM): Tensor[T] {.inline.} =

Possibilities

1. Always use var

   • 👍 Makes sure there is no extra allocation
   • 👍 Consistency
   • 👎 Must declare result before calling function (and computes the shape for im2col for example)
   • 👎 No chaining

2. Always use return values

   • 👍 Consistency
   • 👍 No call burden
   • 🔢 Sometimes return value is a tuple (i.e. forces to read doc for unwrapping)
     - Can be alleviated with named tuple parameters
     - If tuple, it's not easy to chain functions (but is it needed for nn_primitives?)

3. Return value if single, var parameters otherwise?

   • 👎 Inconsistent
   • The worst of both worlds?

[sherjilozair]

Would both be feasible? Not having to allocate extra memory is a huge plus. Memory issues are the biggest bottleneck in model design, and tensorflow and others are considering in-place operations. If arraymancer can support memory reuse from the getgo, that would be great.

[mratsim]

I'm planning to have the low-level API for layers to only use var params (i.e. convolution, maxpooling, etc) and avoid memory allocation. For activations like relu or sigmoid I already provide a mrelu and msigmoid in-place version.

Future considerations

Sink / Destructors (#150)

Ideally with the future sink/destructors I can provide an optimization when only a single ref is used with a signature like this:

proc sigmoid[TT](v: sink Variable[TT]): Variable[TT] {.inline.}=
  msigmoid(v)
  result = v

Object pool / memory regions (#112)

I have started on an object pool due to the cost of allocating memory on Cuda. Custom allocators are planned with Nim destructors, which would ease the implementation of memory regions/object pools to reuse already allocated memory. See stub implementation here.

[mratsim]

Removing the easy tag.

It seems like I would need 3 levels of API:

1. Very low-level that doesn't allocate at all. Workspace buffers will have to be pass as inputs, Fortran-style. In-place activation functions would be there (name to define).
2. Functional API, similar to PyTorch's. To allow chaining it should allow chaining like x.function(weight, bias). This can be build at the Tensor level, with an overload for Variables.
3. Keras-like high level API, which was started by the domain specific language Neural net domain specific language #219.
   It should be syntactic sugar for initializing weights and biases, inferring their shapes and passing them as parameters i.e. x.function instead of x.function(weight, bias)

1 is key for efficient RNNs, as we can reuse the buffers across timesteps instead of allocating a new one every time. This is also required for OpenMP or MPI parallelism at a higher level (model parallelism) as allocations within an OpenMP context will crash.