This is an automated email from the ASF dual-hosted git repository. liuyizhi pushed a commit to branch v1.2.0-java in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git
The following commit(s) were added to refs/heads/v1.2.0-java by this push: new 4177175 improve NDArrayIter to have Builder and ability to specifying names 4177175 is described below commit 41771756420f0bed0edd8c1369f6010d73ebfa26 Author: Yizhi Liu <yizhi...@amazon.com> AuthorDate: Fri Jun 8 17:39:31 2018 -0700 improve NDArrayIter to have Builder and ability to specifying names --- .../scala/org/apache/mxnet/io/NDArrayIter.scala | 90 +++++++++++++++------- .../src/test/scala/org/apache/mxnet/IOSuite.scala | 8 +- 2 files changed, 70 insertions(+), 28 deletions(-) diff --git a/scala-package/core/src/main/scala/org/apache/mxnet/io/NDArrayIter.scala b/scala-package/core/src/main/scala/org/apache/mxnet/io/NDArrayIter.scala index 5108938..ed3c5ad 100644 --- a/scala-package/core/src/main/scala/org/apache/mxnet/io/NDArrayIter.scala +++ b/scala-package/core/src/main/scala/org/apache/mxnet/io/NDArrayIter.scala @@ -23,6 +23,7 @@ import org.apache.mxnet.Base._ import org.apache.mxnet._ import org.slf4j.LoggerFactory +import scala.annotation.varargs import scala.collection.immutable.ListMap /** @@ -38,15 +39,23 @@ import scala.collection.immutable.ListMap * the size of data does not match batch_size. Roll over is intended * for training and can cause problems if used for prediction. */ -class NDArrayIter (data: IndexedSeq[NDArray], label: IndexedSeq[NDArray] = IndexedSeq.empty, - private val dataBatchSize: Int = 1, shuffle: Boolean = false, - lastBatchHandle: String = "pad", - dataName: String = "data", labelName: String = "label") extends DataIter { - private val logger = LoggerFactory.getLogger(classOf[NDArrayIter]) +class NDArrayIter(data: IndexedSeq[(String, NDArray)], + label: IndexedSeq[(String, NDArray)], + private val dataBatchSize: Int, shuffle: Boolean, + lastBatchHandle: String) extends DataIter { + + def this(data: IndexedSeq[NDArray], label: IndexedSeq[NDArray] = IndexedSeq.empty, + dataBatchSize: Int = 1, shuffle: Boolean = false, + lastBatchHandle: String = "pad", + dataName: String = "data", labelName: String = "label") { + this(IO.initData(data, allowEmpty = false, dataName), + IO.initData(label, allowEmpty = true, labelName), + dataBatchSize, shuffle, lastBatchHandle) + } + private val logger = LoggerFactory.getLogger(classOf[NDArrayIter]) - private val (_dataList: IndexedSeq[NDArray], - _labelList: IndexedSeq[NDArray]) = { + val (initData: IndexedSeq[(String, NDArray)], initLabel: IndexedSeq[(String, NDArray)]) = { // data should not be null and size > 0 require(data != null && data.size > 0, "data should not be null and data.size should not be zero") @@ -55,17 +64,17 @@ class NDArrayIter (data: IndexedSeq[NDArray], label: IndexedSeq[NDArray] = Index "label should not be null. Use IndexedSeq.empty if there are no labels") // shuffle is not supported currently - require(shuffle == false, "shuffle is not supported currently") + require(!shuffle, "shuffle is not supported currently") // discard final part if lastBatchHandle equals discard if (lastBatchHandle.equals("discard")) { - val dataSize = data(0).shape(0) + val dataSize = data(0)._2.shape(0) require(dataBatchSize <= dataSize, "batch_size need to be smaller than data size when not padding.") val keepSize = dataSize - dataSize % dataBatchSize - val dataList = data.map(ndArray => {ndArray.slice(0, keepSize)}) + val dataList = data.map { case (name, ndArray) => (name, ndArray.slice(0, keepSize)) } if (!label.isEmpty) { - val labelList = label.map(ndArray => {ndArray.slice(0, keepSize)}) + val labelList = label.map { case (name, ndArray) => (name, ndArray.slice(0, keepSize)) } (dataList, labelList) } else { (dataList, label) @@ -75,13 +84,9 @@ class NDArrayIter (data: IndexedSeq[NDArray], label: IndexedSeq[NDArray] = Index } } - - val initData: IndexedSeq[(String, NDArray)] = IO.initData(_dataList, false, dataName) - val initLabel: IndexedSeq[(String, NDArray)] = IO.initData(_labelList, true, labelName) - val numData = _dataList(0).shape(0) - val numSource = initData.size - var cursor = -dataBatchSize - + val numData = initData(0)._2.shape(0) + val numSource: MXUint = initData.size + private var cursor = -dataBatchSize private val (_provideData: ListMap[String, Shape], _provideLabel: ListMap[String, Shape]) = { @@ -112,8 +117,8 @@ class NDArrayIter (data: IndexedSeq[NDArray], label: IndexedSeq[NDArray] = Index * reset the iterator */ override def reset(): Unit = { - if (lastBatchHandle.equals("roll_over") && cursor>numData) { - cursor = -dataBatchSize + (cursor%numData)%dataBatchSize + if (lastBatchHandle.equals("roll_over") && cursor > numData) { + cursor = -dataBatchSize + (cursor%numData) % dataBatchSize } else { cursor = -dataBatchSize } @@ -154,16 +159,16 @@ class NDArrayIter (data: IndexedSeq[NDArray], label: IndexedSeq[NDArray] = Index newArray } - private def _getData(data: IndexedSeq[NDArray]): IndexedSeq[NDArray] = { + private def _getData(data: IndexedSeq[(String, NDArray)]): IndexedSeq[NDArray] = { require(cursor < numData, "DataIter needs reset.") if (data == null) { null } else { if (cursor + dataBatchSize <= numData) { - data.map(ndArray => {ndArray.slice(cursor, cursor + dataBatchSize)}).toIndexedSeq + data.map { case (_, ndArray) => ndArray.slice(cursor, cursor + dataBatchSize) } } else { // padding - data.map(_padData).toIndexedSeq + data.map { case (_, ndArray) => _padData(ndArray) } } } } @@ -173,7 +178,7 @@ class NDArrayIter (data: IndexedSeq[NDArray], label: IndexedSeq[NDArray] = Index * @return the data of current batch */ override def getData(): IndexedSeq[NDArray] = { - _getData(_dataList) + _getData(initData) } /** @@ -181,7 +186,7 @@ class NDArrayIter (data: IndexedSeq[NDArray], label: IndexedSeq[NDArray] = Index * @return the label of current batch */ override def getLabel(): IndexedSeq[NDArray] = { - _getData(_labelList) + _getData(initLabel) } /** @@ -189,7 +194,7 @@ class NDArrayIter (data: IndexedSeq[NDArray], label: IndexedSeq[NDArray] = Index * @return */ override def getIndex(): IndexedSeq[Long] = { - (cursor.toLong to (cursor + dataBatchSize).toLong).toIndexedSeq + cursor.toLong to (cursor + dataBatchSize).toLong } /** @@ -213,3 +218,36 @@ class NDArrayIter (data: IndexedSeq[NDArray], label: IndexedSeq[NDArray] = Index override def batchSize: Int = dataBatchSize } + +object NDArrayIter { + class Builder() { + private var data: IndexedSeq[(String, NDArray)] = IndexedSeq.empty + private var label: IndexedSeq[(String, NDArray)] = IndexedSeq.empty + private var dataBatchSize: Int = 1 + private var lastBatchHandle: String = "pad" + + def addData(name: String, data: NDArray): Builder = { + this.data = this.data ++ IndexedSeq((name, data)) + this + } + + def addLabel(name: String, label: NDArray): Builder = { + this.label = this.label ++ IndexedSeq((name, label)) + this + } + + def setBatchSize(batchSize: Int): Builder = { + this.dataBatchSize = batchSize + this + } + + def setLastBatchHandle(lastBatchHandle: String): Builder = { + this.lastBatchHandle = lastBatchHandle + this + } + + def build(): NDArrayIter = { + new NDArrayIter(data, label, dataBatchSize, false, lastBatchHandle) + } + } +} diff --git a/scala-package/core/src/test/scala/org/apache/mxnet/IOSuite.scala b/scala-package/core/src/test/scala/org/apache/mxnet/IOSuite.scala index 0f4b7c0..1b922b3 100644 --- a/scala-package/core/src/test/scala/org/apache/mxnet/IOSuite.scala +++ b/scala-package/core/src/test/scala/org/apache/mxnet/IOSuite.scala @@ -24,7 +24,7 @@ import scala.sys.process._ class IOSuite extends FunSuite with BeforeAndAfterAll { - private var tu = new TestUtil + private val tu = new TestUtil test("test MNISTIter & MNISTPack") { // get data @@ -258,7 +258,11 @@ class IOSuite extends FunSuite with BeforeAndAfterAll { assert(batchCount === nBatch0) // test discard - val dataIter1 = new NDArrayIter(data, label, 128, false, "discard") + val dataIter1 = new NDArrayIter.Builder() + .addData("data0", data(0)).addData("data1", data(1)) + .addLabel("label", label(0)) + .setBatchSize(128) + .setLastBatchHandle("discard").build() val nBatch1 = 7 batchCount = 0 while(dataIter1.hasNext) { -- To stop receiving notification emails like this one, please contact liuyi...@apache.org.