<commit_before>package hex.deepwater;

import water.Futures;
import water.H2O;
import water.Job;
import water.MRTask;
import water.gpu.ImageTrain;
import water.parser.BufferedString;
import water.util.Log;
import water.util.PrettyPrint;
import water.util.RandomUtils;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Random;

public class DeepWaterTask extends MRTask<DeepWaterTask> {
  final private boolean _training;
  private DeepWaterModelInfo _localmodel; //per-node state (to be reduced)
  private DeepWaterModelInfo _sharedmodel; //input/output
  int _chunk_node_count = 1;
  float _useFraction;
  boolean _shuffle;
  final Job _job;

  /**
   * Accessor to the object containing the (final) state of the Deep Learning model
   * Should only be queried after calling this.doAll(Frame training)
   * @return "The" final model after one Map/Reduce iteration
   */
  final public DeepWaterModelInfo model_info() {
    assert(_sharedmodel != null);
    return _sharedmodel;
  }

  /**
   * The only constructor
   * @param inputModel Initial model state
   * @param fraction Fraction of rows of the training to train with
   */
  public DeepWaterTask(DeepWaterModelInfo inputModel, float fraction, Job job) {
    _training=true;
    _sharedmodel = inputModel;
    _useFraction=fraction;
    _shuffle = model_info().get_params()._shuffle_training_data;
    _job = job;
  }

  /**
   * Transfer ownership from global (shared) model to local model which will be worked on
   */
  @Override protected void setupLocal(){
    assert(_localmodel == null);
    super.setupLocal();
    _localmodel = _sharedmodel;
    _sharedmodel = null;
    _localmodel.set_processed_local(0);
    final int weightIdx =_fr.find(_localmodel.get_params()._weights_column);
    final int respIdx =_fr.find(_localmodel.get_params()._response_column);
    final int batch_size = _localmodel.get_params()._mini_batch_size;

    // single-threaded logic
    BufferedString bs = new BufferedString();
    int width = 224;
    int height = 224;

    if (_fr.numRows()>Integer.MAX_VALUE) {
      throw H2O.unimpl("Need to implement batching into int-sized chunks.");
    }

    // loop over all images on this node
    ArrayList<Float> train_labels = new ArrayList<>();
    ArrayList<String> train_data = new ArrayList<>();
    for (int i=0; i<_fr.vec(0).length(); ++i) {
      double weight = weightIdx == -1 ? 1 : _fr.vec(weightIdx).at(i);
      if (weight == 0)
        continue;
      String file = _fr.vec(0).atStr(bs, i).toString();
      float response = (float) _fr.vec(respIdx).at(i);
      train_data.add(file);
      train_labels.add(response);
    }
    long seed = 0xDECAF + 0xD00D * _localmodel.get_processed_global();
    Random rng = RandomUtils.getRNG(seed);
    if (_localmodel.get_params()._shuffle_training_data) {
      Collections.shuffle(train_labels, rng);
      rng.setSeed(seed);
      Collections.shuffle(train_data, rng);
    }
    // randomly add more rows to fill up to a multiple of batch_size
    while (train_data.size()%batch_size!=0) {
      int pick = rng.nextInt(train_data.size());
      train_data.add(train_data.get(pick));
      train_labels.add(train_labels.get(pick));
    }
    try {
      long start = System.currentTimeMillis();
      DeepWaterImageIterator img_iter = new DeepWaterImageIterator(train_data, train_labels, batch_size, width, height);
      long end = System.currentTimeMillis();
      Log.info("Time to make Iter: " + PrettyPrint.msecs(end-start, true));

      start = System.currentTimeMillis();
      long nativetime =0;
      Futures fs = new Futures();
      NativeTrainTask ntt = null;
      while(img_iter.Next(fs) && !_job.isStopping()) {
        if (ntt!=null) nativetime +=ntt._timeInMillis;
        float[] data = img_iter.getData();
        float[] labels = img_iter.getLabel();
        long n = _localmodel.get_processed_total();
        Log.info("Trained " + n + " samples. Training on " + Arrays.toString(img_iter.getFiles()));
        _localmodel._imageTrain.setLR(_localmodel.get_params().rate((double)n));
        _localmodel._imageTrain.setMomentum(_localmodel.get_params().momentum((double)n));
        //fork off GPU work, but let the iterator.Next() wait on completion before swapping again
        ntt = new NativeTrainTask(_localmodel._imageTrain, data, labels);
        fs.add(H2O.submitTask(ntt));
        _localmodel.add_processed_local(batch_size);
      }
      nativetime +=ntt._timeInMillis;
      end = System.currentTimeMillis();
      long cputime=end-start- nativetime;
      Log.info("Time for one epoch: " + PrettyPrint.msecs(end-start, true));
      Log.info("Time for data preparation: " + PrettyPrint.msecs(cputime, true));
      Log.info("Time for Native training : " + PrettyPrint.msecs(nativetime, true));

    } catch (IOException e) {
      e.printStackTrace();
    }
  }

  static private class NativeTrainTask extends H2O.H2OCountedCompleter<NativeTrainTask> {
    NativeTrainTask(ImageTrain it, float[] data, float[] labels) {
      _it = it;
      _data = data;
      _labels = labels;
    }
    long _timeInMillis;
    final ImageTrain _it;
    float[] _data;
    float[] _labels;

    @Override
    public void compute2() {
      long start = System.currentTimeMillis();
      _it.train(_data,_labels); //ignore predictions
      long end = System.currentTimeMillis();
      _timeInMillis += end-start;
      tryComplete();<commit_msg>Clean up timing reports.<commit_after>package hex.deepwater;

import water.Futures;
import water.H2O;
import water.Job;
import water.MRTask;
import water.gpu.ImageTrain;
import water.parser.BufferedString;
import water.util.Log;
import water.util.PrettyPrint;
import water.util.RandomUtils;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Random;

public class DeepWaterTask extends MRTask<DeepWaterTask> {
  final private boolean _training;
  private DeepWaterModelInfo _localmodel; //per-node state (to be reduced)
  private DeepWaterModelInfo _sharedmodel; //input/output
  int _chunk_node_count = 1;
  float _useFraction;
  boolean _shuffle;
  final Job _job;

  /**
   * Accessor to the object containing the (final) state of the Deep Learning model
   * Should only be queried after calling this.doAll(Frame training)
   * @return "The" final model after one Map/Reduce iteration
   */
  final public DeepWaterModelInfo model_info() {
    assert(_sharedmodel != null);
    return _sharedmodel;
  }

  /**
   * The only constructor
   * @param inputModel Initial model state
   * @param fraction Fraction of rows of the training to train with
   */
  public DeepWaterTask(DeepWaterModelInfo inputModel, float fraction, Job job) {
    _training=true;
    _sharedmodel = inputModel;
    _useFraction=fraction;
    _shuffle = model_info().get_params()._shuffle_training_data;
    _job = job;
  }

  /**
   * Transfer ownership from global (shared) model to local model which will be worked on
   */
  @Override protected void setupLocal(){
    assert(_localmodel == null);
    super.setupLocal();
    _localmodel = _sharedmodel;
    _sharedmodel = null;
    _localmodel.set_processed_local(0);
    final int weightIdx =_fr.find(_localmodel.get_params()._weights_column);
    final int respIdx =_fr.find(_localmodel.get_params()._response_column);
    final int batch_size = _localmodel.get_params()._mini_batch_size;

    // single-threaded logic
    BufferedString bs = new BufferedString();
    int width = 224;
    int height = 224;

    if (_fr.numRows()>Integer.MAX_VALUE) {
      throw H2O.unimpl("Need to implement batching into int-sized chunks.");
    }

    // loop over all images on this node
    ArrayList<Float> train_labels = new ArrayList<>();
    ArrayList<String> train_data = new ArrayList<>();
    for (int i=0; i<_fr.vec(0).length(); ++i) {
      double weight = weightIdx == -1 ? 1 : _fr.vec(weightIdx).at(i);
      if (weight == 0)
        continue;
      String file = _fr.vec(0).atStr(bs, i).toString();
      float response = (float) _fr.vec(respIdx).at(i);
      train_data.add(file);
      train_labels.add(response);
    }
    long seed = 0xDECAF + 0xD00D * _localmodel.get_processed_global();
    Random rng = RandomUtils.getRNG(seed);
    if (_localmodel.get_params()._shuffle_training_data) {
      Collections.shuffle(train_labels, rng);
      rng.setSeed(seed);
      Collections.shuffle(train_data, rng);
    }
    // randomly add more rows to fill up to a multiple of batch_size
    while (train_data.size()%batch_size!=0) {
      int pick = rng.nextInt(train_data.size());
      train_data.add(train_data.get(pick));
      train_labels.add(train_labels.get(pick));
    }
    try {
      long start = System.currentTimeMillis();
      DeepWaterImageIterator img_iter = new DeepWaterImageIterator(train_data, train_labels, batch_size, width, height);
      long end = System.currentTimeMillis();
      Log.info("Time to make Iter: " + PrettyPrint.msecs(end-start, true));

      start = System.currentTimeMillis();
      long nativetime =0;
      Futures fs = new Futures();
      NativeTrainTask ntt = null;
      while(img_iter.Next(fs) && !_job.isStopping()) {
        if (ntt!=null) nativetime +=ntt._timeInMillis;
        float[] data = img_iter.getData();
        float[] labels = img_iter.getLabel();
        long n = _localmodel.get_processed_total();
        Log.info("Trained " + n + " samples. Training on " + Arrays.toString(img_iter.getFiles()));
        _localmodel._imageTrain.setLR(_localmodel.get_params().rate((double)n));
        _localmodel._imageTrain.setMomentum(_localmodel.get_params().momentum((double)n));
        //fork off GPU work, but let the iterator.Next() wait on completion before swapping again
        ntt = new NativeTrainTask(_localmodel._imageTrain, data, labels);
        fs.add(H2O.submitTask(ntt));
        _localmodel.add_processed_local(batch_size);
      }
      nativetime +=ntt._timeInMillis;
      end = System.currentTimeMillis();
      Log.info("Time for one epoch: " + PrettyPrint.msecs(end-start, true));
      Log.info("Time for Native training : " + PrettyPrint.msecs(nativetime, true));

    } catch (IOException e) {
      e.printStackTrace();
    }
  }

  static private class NativeTrainTask extends H2O.H2OCountedCompleter<NativeTrainTask> {
    NativeTrainTask(ImageTrain it, float[] data, float[] labels) {
      _it = it;
      _data = data;
      _labels = labels;
    }
    long _timeInMillis;
    final ImageTrain _it;
    float[] _data;
    float[] _labels;

    @Override
    public void compute2() {
      long start = System.currentTimeMillis();
      _it.train(_data,_labels); //ignore predictions
      long end = System.currentTimeMillis();
      _timeInMillis += end-start;
      tryComplete();
{'fraction_non_alphanumeric': 0.11274464761554719, 'fraction_numerical': 0.004778094275475512, 'mean_word_length': 3.4679802955665027, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, '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': '4820322', 'n_tokens_mistral': 3405, 'n_tokens_neox': 3344, 'n_words': 823}