Module API

The module API provides an intermediate and high-level interface for performing computation with neural networks in MXNet. A module is an instance of subclasses of the BaseModule. The most widely used module class is called Module. Module wraps a Symbol and one or more Executors. For a full list of functions, see BaseModule. A subclass of modules might have extra interface functions. This topic provides some examples of common use cases. All of the module APIs are in the Module namespace.

Preparing a Module for Computation

To construct a module, refer to the constructors for the module class. For example, the Module class accepts a Symbol as input:

    import org.apache.mxnet._
    import org.apache.mxnet.module.{FitParams, Module}

    // construct a simple MLP
    val data = Symbol.Variable("data")
    val fc1 = Symbol.api.FullyConnected(Some(data), num_hidden = 128, name = "fc1")
    val act1 = Symbol.api.Activation(Some(fc1), "relu", "relu1")
    val fc2 = Symbol.api.FullyConnected(Some(act1), num_hidden = 64, name = "fc2")
    val act2 = Symbol.api.Activation(Some(fc2), "relu", "relu2")
    val fc3 = Symbol.api.FullyConnected(Some(act2), num_hidden = 10, name = "fc3")
    val out = Symbol.api.SoftmaxOutput(fc3, name = "softmax")

    // construct the module
    val mod = new Module(out)

By default, context is the CPU. If you need data parallelization, you can specify a GPU context or an array of GPU contexts.

Before you can compute with a module, you need to call bind() to allocate the device memory and initParams() or SetParams() to initialize the parameters. If you simply want to fit a module, you don't need to call bind() and initParams() explicitly, because the fit() function automatically calls them if they are needed.

    mod.bind(dataShapes = train_dataiter.provideData, labelShapes = Some(train_dataiter.provideLabel))
    mod.initParams()

Now you can compute with the module using functions like forward(), backward(), etc.

Training, Predicting, and Evaluating

Modules provide high-level APIs for training, predicting, and evaluating. To fit a module, call the fit() function with some DataIters:

    import org.apache.mxnet.optimizer.SGD
    val mod = new Module(softmax)

    mod.fit(train_dataiter, evalData = scala.Option(eval_dataiter), \
    numEpoch = n_epoch, fitParams = new FitParams()\
    .setOptimizer(new SGD(learningRate = 0.1f, momentum = 0.9f, wd = 0.0001f)))

The interface is very similar to the old FeedForward class. You can pass in batch-end callbacks using setBatchEndCallback and epoch-end callbacks using setEpochEndCallback. You can also set parameters using methods like setOptimizer and setEvalMetric. To learn more about the FitParams(), see the API page. To predict with a module, call predict() with a DataIter:

    mod.predict(val_dataiter)

The module collects and returns all of the prediction results. For more details about the format of the return values, see the documentation for the predict() function.

When prediction results might be too large to fit in memory, use the predictEveryBatch API:

    val preds = mod.predictEveryBatch(val_dataiter)
    val_dataiter.reset()
    var i = 0
    while (val_dataiter.hasNext) {
       val batch = val_dataiter.next()
       val predLabel: Array[Int] = NDArray.argmax_channel(preds(i)(0)).toArray.map(_.toInt)
       val label = batch.label(0).toArray.map(_.toInt)
       //do something...
       i += 1
    }

If you need to evaluate on a test set and don't need the prediction output, call the score() function with a DataIter and an EvalMetric:

    mod.score(val_dataiter, metric)

This runs predictions on each batch in the provided DataIter and computes the evaluation score using the provided EvalMetric. The evaluation results are stored in metric so that you can query later.

Saving and Loading Module Parameters

To save the module parameters in each training epoch, use a checkpoint callback:

    val modelPrefix: String = "mymodel"

    for (epoch <- 0 until 5) {
      while(train_dataiter.hasNext){
          // forward backward pass
         //do something...
       }
        val checkpoint = mod.saveCheckpoint(modelPrefix, epoch, saveOptStates = true)

    }

To load the saved module parameters, call the loadCheckpoint function:

    val mod = Module.loadCheckpoint(modelPrefix, loadModelEpoch, loadOptimizerStates = true)

To initialize parameters, Bind the symbols to construct executors first with bind method. Then, initialize the parameters and auxiliary states by calling initParams() method.

    mod.bind(dataShapes = train_dataiter.provideData, labelShapes = Some(train_dataiter.provideLabel))
    mod.initParams()

To get current parameters, use getParams method.

    val (argParams, auxParams) = mod.getParams

To assign parameter and aux state values, use setParams method.

    mod.setParams(argParams, auxParams)

To resume training from a saved checkpoint, instead of calling setParams(), directly call fit(), passing the loaded parameters, so that fit() knows to start from those parameters instead of initializing randomly:

    mod.fit(..., fitParams=new FitParams().setArgParams(argParams).\
    setAuxParams(auxParams).setBeginEpoch(beginEpoch))

Create an object of the FitParams() class, and then use it to call the setBeginEpoch() method to pass beginEpoch so that fit() knows to resume from a saved epoch.

Next Steps

  • See Model API for an alternative simple high-level interface for training neural networks.
  • See Symbolic API for operations on NDArrays that assemble neural networks from layers.
  • See IO Data Loading API for parsing and loading data.
  • See NDArray API for vector/matrix/tensor operations.
  • See KVStore API for multi-GPU and multi-host distributed training.