<commit_before>/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.intel.analytics.bigdl.optim

import com.intel.analytics.bigdl.nn._
import com.intel.analytics.bigdl._
import com.intel.analytics.bigdl.parameters.{AllReduceParameterManager, ImprovedAllReduceParameterManager, OneReduceParameterManager}
import com.intel.analytics.bigdl.tensor.{Storage, Tensor}
import com.intel.analytics.bigdl.utils.{Engine, RandomGenerator, T}
import org.apache.log4j.{Level, Logger}
import org.apache.spark.SparkContext
import org.scalatest.{BeforeAndAfter, BeforeAndAfterAll, FlatSpec, Matchers}

object DropSlowModuleOptimizerSpec {
  val input1: Tensor[Double] = Tensor[Double](Storage[Double](Array(0.0, 1.0, 0.0, 1.0)))
  val output1 = 0.0
  val input2: Tensor[Double] = Tensor[Double](Storage[Double](Array(1.0, 0.0, 1.0, 0.0)))
  val output2 = 1.0
  var plusOne = 0.0
}

class DropSlowModuleOptimizerSpec extends FlatSpec with Matchers with BeforeAndAfter
  with BeforeAndAfterAll {

  Logger.getLogger("org").setLevel(Level.WARN)
  Logger.getLogger("akka").setLevel(Level.WARN)

  import EpochOptimizerSpec._

  var sc: SparkContext = null

  var dataSet: ShuffleBatchDataSet[(Int), Double] = null

  var dataSet2: ShuffleBatchDataSet[(Int), Double] = null

  val (mseModule, mseWeight, mseGradient) = {
    val mlp = new Sequential[Double]
    mlp.add(new Linear(4, 2))
    mlp.add(new Sigmoid)
    mlp.add(new Linear(2, 1))
    mlp.add(new Sigmoid)

    val (weight, grad) = mlp.getParameters()
    (mlp, weight, grad)
  }

  val (crnModule, crnWeight, crnGradient) = {
    val mlp = new Sequential[Double]
    mlp.add(new Linear(4, 2))
    mlp.add(new LogSoftMax)

    val (weight, grad) = mlp.getParameters()
    (mlp, weight, grad)
  }

  before {
    sc = new SparkContext("local[1]", "EpochOptimizerSpec")

    val subModuleNumber = 4
    DropSlowModuleGradAggEpochOptimizer.subModuleNumber = subModuleNumber

    dataSet = new ShuffleBatchDataSet[(Int), Double](
      sc.makeRDD(0 to (256 * 8), 8),
      (seq, input, target) => {
        val size = seq.size
        input.resize(Array(size, 4))
        target.resize(Array(size))
        var i = 0
        while (i < size) {
          if (i % 2 == 0) {
            target.setValue(i + 1, output1 + plusOne)
            input.select(1, i + 1).copy(input1)
          } else {
            target.setValue(i + 1, output2 + plusOne)
            input.select(1, i + 1).copy(input2)
          }
          i += 1
        }
        (input, target)
      }, 4, 4 * subModuleNumber, 1)


    dataSet2 = new ShuffleBatchDataSet[(Int), Double](
      sc.makeRDD(0 to (256 * 8), 8),
      (seq, input, target) => {
        val size = seq.size
        input.resize(Array(size, 4))
        target.resize(Array(size))
        var i = 0
        while (i < size) {
          if (i % 2 == 0) {
            target.setValue(i + 1, output1 + plusOne)
            input.select(1, i + 1).copy(input1)
          } else {
            target.setValue(i + 1, output2 + plusOne)
            input.select(1, i + 1).copy(input2)
          }
          i += 1
        }
        (input, target)
      }, 4 * subModuleNumber,
      4 * subModuleNumber, 1)

    RandomGenerator.RNG.setSeed(1000)
    mseWeight.fill(0.125)
    crnWeight.fill(0.125)
    plusOne = 0.0
    Engine.default.setPoolSize(1)
  }

  after {
    if (sc != null) {
      sc.stop()
    }
  }

  "An Artificial Neural Network with MSE and LBFGS" should "be trained with good result" in {
    val pm = new ImprovedAllReduceParameterManager[Double](mseWeight, dataSet.partitions())
    val optimizer = new DropSlowModuleGradAggEpochOptimizer[Double](mseModule, MSECriterion[Double],
      new LBFGS, pm, dataSet, new Metrics)
    optimizer.setMaxEpoch(5)
    optimizer.optimize()

    mseWeight.copy(optimizer.parameters)
    val result1 = mseModule.forward(input1).toTensor[Double]
    result1(Array(1)) should be(0.0 +- 1e-2)

    val result2 = mseModule.forward(input2).toTensor[Double]
    result2(Array(1)) should be(1.0 +- 1e-2)
  }

  "An Artificial Neural Network with MSE and SGD" should "be trained with good result" in {
    val pm = new ImprovedAllReduceParameterManager[Double](mseWeight, dataSet.partitions())
    val optimizer = new DropSlowModuleGradAggEpochOptimizer[Double](mseModule, MSECriterion[Double],
      new SGD, pm, dataSet, new Metrics, T("learningRate" -> 20.0))
    optimizer.setMaxEpoch(5)
    optimizer.optimize()

    mseWeight.copy(optimizer.parameters)
    val result1 = mseModule.forward(input1).toTensor[Double]
    result1(Array(1)) should be(0.0 +- 5e-2)

    val result2 = mseModule.forward(input2).toTensor[Double]
    result2(Array(1)) should be(1.0 +- 5e-2)
  }

  "An Artificial Neural Network with Cross Entropy and LBFGS" should
    "be trained with good result" in {
    plusOne = 1.0
    val pm = new ImprovedAllReduceParameterManager[Double](crnWeight, dataSet.partitions())
    val optimizer = new DropSlowModuleGradAggEpochOptimizer[Double](crnModule,
      ClassNLLCriterion[Double](), new LBFGS, pm, dataSet, new Metrics)
    optimizer.setMaxEpoch(3)
    optimizer.optimize()

    crnWeight.copy(optimizer.parameters)
    val result1 = crnModule.forward(input1).toTensor[Double]
    result1.max(1)._2(Array(1)) should be(1.0)

    val result2 = crnModule.forward(input2).toTensor[Double]
    result2.max(1)._2(Array(1)) should be(2.0)
  }

  "An Artificial Neural Network with Cross Entropy and SGD" should
    "be trained with good result" in {
    plusOne = 1.0
    val pm = new ImprovedAllReduceParameterManager[Double](crnWeight, dataSet.partitions())
    val optimizer = new DropSlowModuleGradAggEpochOptimizer[Double](crnModule,
      ClassNLLCriterion[Double](), new SGD, pm, dataSet, new Metrics, T("learningRate" -> 20.0))
    optimizer.setMaxEpoch(1)
    optimizer.optimize()

    crnWeight.copy(optimizer.parameters)
    val result1 = crnModule.forward(input1).toTensor[Double]
    result1.max(1)._2(Array(1)) should be(1.0)

    val result2 = crnModule.forward(input2).toTensor[Double]
    result2.max(1)._2(Array(1)) should be(2.0)
  }
}<commit_msg>add more ut for drop module optimizer<commit_after>/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.intel.analytics.bigdl.optim

import com.intel.analytics.bigdl.nn._
import com.intel.analytics.bigdl._
import com.intel.analytics.bigdl.parameters.{AllReduceParameterManager, ImprovedAllReduceParameterManager, OneReduceParameterManager}
import com.intel.analytics.bigdl.tensor.{Storage, Tensor}
import com.intel.analytics.bigdl.utils.{Engine, RandomGenerator, T}
import org.apache.log4j.{Level, Logger}
import org.apache.spark.SparkContext
import org.scalatest.{BeforeAndAfter, BeforeAndAfterAll, FlatSpec, Matchers}

object DropSlowModuleOptimizerSpec {
  val input1: Tensor[Double] = Tensor[Double](Storage[Double](Array(0.0, 1.0, 0.0, 1.0)))
  val output1 = 0.0
  val input2: Tensor[Double] = Tensor[Double](Storage[Double](Array(1.0, 0.0, 1.0, 0.0)))
  val output2 = 1.0
  var plusOne = 0.0
}

class DropSlowModuleOptimizerSpec extends FlatSpec with Matchers with BeforeAndAfter
  with BeforeAndAfterAll {

  Logger.getLogger("org").setLevel(Level.WARN)
  Logger.getLogger("akka").setLevel(Level.WARN)

  import EpochOptimizerSpec._

  var sc: SparkContext = null

  var dataSet: ShuffleBatchDataSet[(Int), Double] = null

  var dataSet2: ShuffleBatchDataSet[(Int), Double] = null

  val (mseModule, mseWeight, mseGradient) = {
    val mlp = new Sequential[Double]
    mlp.add(new Linear(4, 2))
    mlp.add(new Sigmoid)
    mlp.add(new Linear(2, 1))
    mlp.add(new Sigmoid)

    val (weight, grad) = mlp.getParameters()
    (mlp, weight, grad)
  }

  val (crnModule, crnWeight, crnGradient) = {
    val mlp = new Sequential[Double]
    mlp.add(new Linear(4, 2))
    mlp.add(new LogSoftMax)

    val (weight, grad) = mlp.getParameters()
    (mlp, weight, grad)
  }

  before {
    sc = new SparkContext("local[1]", "EpochOptimizerSpec")

    val subModuleNumber = 4
    DropSlowModuleGradAggEpochOptimizer.subModuleNumber = subModuleNumber

    dataSet = new ShuffleBatchDataSet[(Int), Double](
      sc.makeRDD(0 to (256 * 8), 8),
      (seq, input, target) => {
        val size = seq.size
        input.resize(Array(size, 4))
        target.resize(Array(size))
        var i = 0
        while (i < size) {
          if (i % 2 == 0) {
            target.setValue(i + 1, output1 + plusOne)
            input.select(1, i + 1).copy(input1)
          } else {
            target.setValue(i + 1, output2 + plusOne)
            input.select(1, i + 1).copy(input2)
          }
          i += 1
        }
        (input, target)
      }, 4, 4 * subModuleNumber, 1)


    dataSet2 = new ShuffleBatchDataSet[(Int), Double](
      sc.makeRDD(0 to (256 * 8), 8),
      (seq, input, target) => {
        val size = seq.size
        input.resize(Array(size, 4))
        target.resize(Array(size))
        var i = 0
        while (i < size) {
          if (i % 2 == 0) {
            target.setValue(i + 1, output1 + plusOne)
            input.select(1, i + 1).copy(input1)
          } else {
            target.setValue(i + 1, output2 + plusOne)
            input.select(1, i + 1).copy(input2)
          }
          i += 1
        }
        (input, target)
      }, 4 * subModuleNumber,
      4 * subModuleNumber, 1)

    RandomGenerator.RNG.setSeed(1000)
    mseWeight.fill(0.125)
    crnWeight.fill(0.125)
    plusOne = 0.0
    Engine.default.setPoolSize(1)
  }

  after {
    if (sc != null) {
      sc.stop()
    }
  }

  "An Artificial Neural Network with MSE and LBFGS" should "be trained with good result" in {
    val pm = new ImprovedAllReduceParameterManager[Double](mseWeight, dataSet.partitions())
    val optimizer = new DropSlowModuleGradAggEpochOptimizer[Double](mseModule, MSECriterion[Double],
      new LBFGS, pm, dataSet, new Metrics)
    optimizer.setMaxEpoch(5)
    optimizer.optimize()

    mseWeight.copy(optimizer.parameters)
    val result1 = mseModule.forward(input1).toTensor[Double]
    result1(Array(1)) should be(0.0 +- 1e-2)

    val result2 = mseModule.forward(input2).toTensor[Double]
    result2(Array(1)) should be(1.0 +- 1e-2)
  }

//  it should "be same with grad agg optimizer" in {
//    val pm = new AllReduceParameterManager[Double](mseWeight, dataSet.partitions())
//    val optimizer = new BetterGradAggEpochOptimizer[Double](mseModule, MSECriterion[Double],
//      new LBFGS, pm, dataSet, new Metrics, DropSlowModuleGradAggEpochOptimizer.subModuleNumber)
//    optimizer.setMaxEpoch(1)
//    optimizer.optimize()
//
//    val test = pm.getParameter().clone()
//    val pm = new OneReduceParameterManager[Double](mseWeight, dataSet.partitions())
//    val optimizer = new BetterGradAggEpochOptimizer[Double](mseModule, MSECriterion[Double],
//      new LBFGS, pm, dataSet, new Metrics, DropSlowModuleGradAggEpochOptimizer.subModuleNumber)
//    optimizer.setMaxEpoch(1)
//    optimizer.optimize()
//
//    val test = pm.getParameter().clone()
//    mseWeight.fill(0.125)
//    val pm2 = new OneReduceParameterManager[Double](mseWeight, dataSet2.partitions())
//    val optimizer2 = new GradAggEpochOptimizer[Double](mseModule, MSECriterion[Double],
//      new LBFGS, pm2, dataSet2, new Metrics)
//    optimizer2.setMaxEpoch(1)
//    optimizer2.optimize()
//
//    println(test)
//    println(mseWeight)
//
//    test.map(mseWeight, (a, b) => {
//      a should be(b)
//      a
//    })
//  }

//  it should "be same with grad agg optimizer with multi stack dataset" in {
//    val pm = new ImprovedAllReduceParameterManager[Double](mseWeight, dataSet.partitions())
//    val optimizer = new DropSlowModuleGradAggEpochOptimizer[Double](mseModule,
//    MSECriterion[Double], new LBFGS, pm, dataSet, new Metrics)
//    optimizer.setMaxEpoch(1)
//    optimizer.optimize()
//
//    val test = optimizer.parameters
//    mseWeight.fill(0.125)
//
//    val pm3 = new OneReduceParameterManager[Double](mseWeight, dataSet.partitions())
//    val optimizer3 = new BetterGradAggEpochOptimizer[Double](mseModule, MSECriterion[Double],
//      new LBFGS, pm3, dataSet, new Metrics, DropSlowModuleGradAggEpochOptimizer.subModuleNumber)
//    optimizer3.setMaxEpoch(1)
//    optimizer3.optimize()
//
//    val test3 = pm3.getParameter().clone()
//
// //    val pm2 = new OneReduceParameterManager[Double](mseWeight, dataSet.partitions())
// //    val optimizer2 = new GradAggEpochOptimizer[Double](mseModule, MSECriterion[Double],
// //      new LBFGS, pm2, dataSet, new Metrics)
// //    optimizer2.setMaxEpoch(1)
// //    optimizer2.optimize()
//
//    println(test)
//    println(test3)
// //    println(mseWeight)
//
// //    test.map(mseWeight, (a, b) => {
//    test.map(test3, (a, b) => {
//      a should be(b)
//      a
//    })
//  }

  "An Artificial Neural Network with MSE and SGD" should "be trained with good result" in {
    val pm = new ImprovedAllReduceParameterManager[Double](mseWeight, dataSet.partitions())
    val optimizer = new DropSlowModuleGradAggEpochOptimizer[Double](mseModule, MSECriterion[Double],
      new SGD, pm, dataSet, new Metrics, T("learningRate" -> 20.0))
    optimizer.setMaxEpoch(5)
    optimizer.optimize()

    mseWeight.copy(optimizer.parameters)
    val result1 = mseModule.forward(input1).toTensor[Double]
    result1(Array(1)) should be(0.0 +- 5e-2)

    val result2 = mseModule.forward(input2).toTensor[Double]
    result2(Array(1)) should be(1.0 +- 5e-2)
  }

  it should "be same with grad agg optimizer" in {
    val pm = new ImprovedAllReduceParameterManager[Double](mseWeight, dataSet.partitions())
    val optimizer = new DropSlowModuleGradAggEpochOptimizer[Double](mseModule, MSECriterion[Double],
      new SGD, pm, dataSet, new Metrics, T("learningRate" -> 20.0))
    optimizer.setMaxEpoch(1)
    optimizer.optimize()

    val test = optimizer.parameters

    mseWeight.fill(0.125)
    val pm2 = new OneReduceParameterManager[Double](mseWeight, dataSet2.partitions())
    val optimizer2 = new GradAggEpochOptimizer[Double](mseModule, MSECriterion[Double],
      new SGD, pm2, dataSet2, new Metrics, T("learningRate" -> 20.0))
    optimizer2.setMaxEpoch(1)
    optimizer2.optimize()

    test.map(mseWeight, (a, b) => {
      a should be(b)
      a
    })
  }

  "An Artificial Neural Network with Cross Entropy and LBFGS" should
    "be trained with good result" in {
    plusOne = 1.0
    val pm = new ImprovedAllReduceParameterManager[Double](crnWeight, dataSet.partitions())
    val optimizer = new DropSlowModuleGradAggEpochOptimizer[Double](crnModule,
      ClassNLLCriterion[Double](), new LBFGS, pm, dataSet, new Metrics)
    optimizer.setMaxEpoch(3)
    optimizer.optimize()

    crnWeight.copy(optimizer.parameters)
    val result1 = crnModule.forward(input1).toTensor[Double]
    result1.max(1)._2(Array(1)) should be(1.0)

    val result2 = crnModule.forward(input2).toTensor[Double]
    result2.max(1)._2(Array(1)) should be(2.0)
  }

  it should "be same with grad agg optimizer" in {
    plusOne = 1.0
    val pm = new ImprovedAllReduceParameterManager[Double](crnWeight, dataSet.partitions())
    val optimizer = new DropSlowModuleGradAggEpochOptimizer[Double](crnModule,
      ClassNLLCriterion[Double](), new LBFGS, pm, dataSet, new Metrics)
    optimizer.setMaxEpoch(3)
    optimizer.optimize()
    val test = optimizer.parameters

    crnWeight.fill(0.125)
    val pm3 = new AllReduceParameterManager[Double](crnWeight, dataSet.partitions())
    val optimizer3 = new BetterGradAggEpochOptimizer[Double](crnModule, ClassNLLCriterion[Double](),
      new LBFGS, pm3, dataSet, new Metrics, DropSlowModuleGradAggEpochOptimizer.subModuleNumber)
    optimizer3.setMaxEpoch(3)
    optimizer3.optimize()
    val test3 = pm3.getParameter()

    test.map(test3, (a, b) => {
      a should be(b)
      a
    })
  }

  "An Artificial Neural Network with Cross Entropy and SGD" should
    "be trained with good result" in {
    plusOne = 1.0
    val pm = new ImprovedAllReduceParameterManager[Double](crnWeight, dataSet.partitions())
    val optimizer = new DropSlowModuleGradAggEpochOptimizer[Double](crnModule,
      ClassNLLCriterion[Double](), new SGD, pm, dataSet, new Metrics, T("learningRate" -> 20.0))
    optimizer.setMaxEpoch(1)
    optimizer.optimize()

    crnWeight.copy(optimizer.parameters)
    val result1 = crnModule.forward(input1).toTensor[Double]
    result1.max(1)._2(Array(1)) should be(1.0)

    val result2 = crnModule.forward(input2).toTensor[Double]
    result2.max(1)._2(Array(1)) should be(2.0)
  }

  it should "be same with grad agg optimizer" in {
    plusOne = 1.0
    val pm = new ImprovedAllReduceParameterManager[Double](crnWeight, dataSet.partitions())
    val optimizer = new DropSlowModuleGradAggEpochOptimizer[Double](crnModule,
      ClassNLLCriterion[Double](), new SGD, pm, dataSet, new Metrics, T("learningRate" -> 20.0))
    optimizer.setMaxEpoch(1)
    optimizer.optimize()

    val test = optimizer.parameters

    crnWeight.fill(0.125)
    val pm3 = new AllReduceParameterManager[Double](crnWeight, dataSet.partitions())
    val optimizer3 = new BetterGradAggEpochOptimizer[Double](crnModule, ClassNLLCriterion[Double](),
      new SGD, pm3, dataSet, new Metrics, DropSlowModuleGradAggEpochOptimizer.subModuleNumber,
      T("learningRate" -> 20.0))
    optimizer3.setMaxEpoch(1)
    optimizer3.optimize()
    val test3 = pm3.getParameter().clone()
//    crnWeight.fill(0.125)
//    val pm2 = new OneReduceParameterManager[Double](crnWeight, dataSet2.partitions())
//    val optimizer2 = new GradAggEpochOptimizer[Double](crnModule, ClassNLLCriterion[Double](),
//      new SGD, pm2, dataSet2, new Metrics, T("learningRate" -> 20.0))
//    optimizer2.setMaxEpoch(1)
//    optimizer2.optimize()

//    test3.map(crnWeight, (a, b) => {
      test.map(test3, (a, b) => {
      a should be(b)
      a
    })
  }
}
{'fraction_non_alphanumeric': 0.12069588560218884, 'fraction_numerical': 0.024366320788808012, 'mean_word_length': 3.3877689694224236, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 23, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2501230', 'n_tokens_mistral': 6727, 'n_tokens_neox': 6445, 'n_words': 1507}