Base structure for Native Operators

examples/julia-softmax.jl

@@ -0,0 +1,63 @@
+using MXNet
+
+import MXNet.mx.Native: Operator, list_arguments, list_outputs, infer_shape,
+                        forward, backward, need_top_grad, create_op
+type JuliaSoftmax <: Operator end
+
+(op::JuliaSoftmax)(;kwargs...) = create_op(op;kwargs...)
+
+list_arguments(:: JuliaSoftmax) = ["data", "label"]
+list_outputs(:: JuliaSoftmax) = ["output"]
+need_top_grad(:: JuliaSoftmax) = false
+
+function infer_shape(::JuliaSoftmax, in_shapes :: Vector{Vector{UInt32}})
+  data_shape = in_shapes[1]
+  label_shape = [last(data_shape)]
+  output_shape = data_shape
+  return (data_shape, label_shape), (output_shape, )
+end
+
+function forward(::JuliaSoftmax, in_data :: Vector{mx.NDArray}, out_data :: Vector{mx.NDArray})
+  info("Entering forward")
+  x = in_data[1]
+  y = out_data[1]
+
+  @mx.nd_as_jl ro=x rw=y begin
+    y[:] = exp(x - maximum(x, 1))
+    y /= sum(y, 1)
+  end
+  info("Leaving forward")
+end
+
+#TODO: Correct gradient
+function backward(::JuliaSoftmax, out_grad :: Vector{mx.NDArray}, in_data :: Vector{mx.NDArray}, out_data :: Vector{mx.NDArray}, in_grad :: Vector{mx.NDArray})
+  info("Entering backward")
+  label = in_data[2]
+  y = out_data[1]
+  dx = in_grad[1]
+
+  @mx.nd_as_jl ro=(label, y) rw=dx begin
+    dx[:] = y
+  end
+  info("Leaving backward")
+end
+
+#define mlp
+data = mx.Variable("data")
+fc1 = mx.FullyConnected(data = data, name="fc1", num_hidden=128)
+act1 = mx.Activation(data = fc1, name="relu1", act_type="relu")
+fc2 = mx.FullyConnected(data = act1, name="fc2", num_hidden=64)
+act2 = mx.Activation(data = fc2, name="relu2", act_type="relu")
+fc3 = mx.FullyConnected(data = act2, name="fc3", num_hidden=10)
+
+# Setup Native operator
+mysoftmax = JuliaSoftmax()
+mlp = mysoftmax(name = "softmax", data=fc3)
+
+model = mx.FeedForward(mlp, context = mx.cpu())
+optimizer = mx.SGD(lr = 0.1, momentum = 0.9, weight_decay = 0.00001)
+
+include("mnist/mnist-data.jl")
+train_provider, eval_provider = get_mnist_providers(100)
+mx.fit(model, optimizer, train_provider, eval_data=eval_provider, n_epoch =20 )
+

src/MXNet.jl

@@ -37,6 +37,7 @@ include("kvstore.jl")
 include("callback.jl")
 include("model.jl")
 
+include("nativeops.jl")
 include("visualize.jl")
 
 include("nn-factory.jl")

src/base.jl

@@ -60,6 +60,21 @@ macro mxcall(fv, argtypes, args...)
   end
 end
 
+"Utility macro to call MXNet API functions from a threadpool"
+macro mxthreadcall(fv, argtypes, args...)
+  f = eval(fv)
+  args = map(esc, args)
+  quote
+    _mxret = @threadcall( ($(Meta.quot(f)), $MXNET_LIB),
+                    Cint, $argtypes, $(args...) )
+    if _mxret != 0
+      err_msg = mx_get_last_error()
+      throw(MXError(err_msg))
+    end
+  end
+end
+
+
 ################################################################################
 # Handle types
 ################################################################################

src/executor.jl

@@ -164,7 +164,9 @@ function forward(self :: Executor; is_train::Bool=false, kwargs...)
     copy!(self.arg_dict[k], v)
   end
 
-  @mxcall(:MXExecutorForward, (MX_handle, Cint), self, is_train)
+  println("Forward")
+  @mxthreadcall(:MXExecutorForward, (MX_handle, Cint), self, is_train)
+  println("Forward")
 end
 
 function backward(self :: Executor)
@@ -175,7 +177,9 @@ function backward(self :: Executor, out_grad :: NDArray)
 end
 function backward(self :: Executor, out_grads :: Vector{NDArray})
   out_grads = MX_handle[out_grads...]
-  @mxcall(:MXExecutorBackward, (MX_handle, MX_uint, Ptr{MX_handle}), self, length(out_grads), out_grads)
+  println("Backward")
+  @mxthreadcall(:MXExecutorBackward, (MX_handle, MX_uint, Ptr{MX_handle}), self, length(out_grads), out_grads)
+  println("Backward")
 end