Train a Deep Learning model, assumes that all members are populated
If checkpoint == null, then start training a new model, otherwise continue from a checkpoint
188688
private void buildModel() {
    DeepLearningModel cp = null;
    if (checkpoint == null) {
      cp = initModel();
      cp.start_training(null);
    } else {
      final DeepLearningModel previous = UKV.get(checkpoint);
      if (previous == null) throw new IllegalArgumentException("Checkpoint not found.");
      Log.info("Resuming from checkpoint.");
      if (n_folds != 0) {
        throw new UnsupportedOperationException("n_folds must be 0: Cross-validation is not supported during checkpoint restarts.");
      }
      else {
        ((ValidatedJob)previous.job()).xval_models = null; //remove existing cross-validation keys after checkpoint restart
      }
      if (source == null || (previous.model_info().get_params().source != null && !Arrays.equals(source._key._kb, previous.model_info().get_params().source._key._kb))) {
        throw new IllegalArgumentException("source must be the same as for the checkpointed model.");
      }
      autoencoder = previous.model_info().get_params().autoencoder;
      if (!autoencoder && (response == null || !source.names()[source.find(response)].equals(previous.responseName()))) {
        throw new IllegalArgumentException("response must be the same as for the checkpointed model.");
      }
//      if (!autoencoder && (response == null || !Arrays.equals(response._key._kb, previous.model_info().get_params().response._key._kb))) {
//        throw new IllegalArgumentException("response must be the same as for the checkpointed model.");
//      }
      if (Utils.difference(ignored_cols, previous.model_info().get_params().ignored_cols).length != 0
              || Utils.difference(previous.model_info().get_params().ignored_cols, ignored_cols).length != 0) {
        ignored_cols = previous.model_info().get_params().ignored_cols;
        Log.warn("Automatically re-using ignored_cols from the checkpointed model.");
      }
      if ((validation == null) == (previous._validationKey != null)
              || (validation != null && validation._key != null && previous._validationKey != null
              && !Arrays.equals(validation._key._kb, previous._validationKey._kb))) {
        throw new IllegalArgumentException("validation must be the same as for the checkpointed model.");
      }
      if (classification != previous.model_info().get_params().classification) {
        Log.warn("Automatically switching to " + ((classification=!classification) ? "classification" : "regression") + " (same as the checkpointed model).");
      }
      epochs += previous.epoch_counter; //add new epochs to existing model
      Log.info("Adding " + String.format("%.3f", previous.epoch_counter) + " epochs from the checkpointed model.");
      try {
        final DataInfo dataInfo = prepareDataInfo();
        cp = new DeepLearningModel(previous, destination_key, job_key, dataInfo);
        cp.write_lock(self());
        cp.start_training(previous);
        assert(state==JobState.RUNNING);
        final DeepLearning A = cp.model_info().get_params();
        Object B = this;
        for (Field fA : A.getClass().getDeclaredFields()) {
          if (Utils.contains(cp_modifiable, fA.getName())) {
            if (!expert_mode && Utils.contains(expert_options, fA.getName())) continue;
            for (Field fB : B.getClass().getDeclaredFields()) {
              if (fA.equals(fB)) {
                try {
                  if (fB.get(B) == null || fA.get(A) == null || !fA.get(A).toString().equals(fB.get(B).toString())) { // if either of the two parameters is null, skip the toString()
                    if (fA.get(A) == null && fB.get(B) == null) continue; //if both parameters are null, we don't need to do anything
                    Log.info("Applying user-requested modification of '" + fA.getName() + "': " + fA.get(A) + " -> " + fB.get(B));
                    fA.set(A, fB.get(B));
                  }
                } catch (IllegalAccessException e) {
                  e.printStackTrace();
                }
              }
            }
          }
        }
        if (A.n_folds != 0) {
          Log.warn("Disabling cross-validation: Not supported when resuming training from a checkpoint.");
          A.n_folds = 0;
        }
        cp.update(self());
      } finally {
        if (cp != null) cp.unlock(self());
      }
    }
    trainModel(cp);
    cp.stop_training();
  }
494352
def train(self, **kwargs):
        
        log.info("Starting to train...")
        if not self.setup_run:
            self.setup()
        self.trained_tasks = []
        for task in self.tasks:
            data = self.reformatted_input[task.data_format]['data']
            target = self.reformatted_input[task.data_format]['target']
            if data is None:
                raise Exception("Data cannot be none.  Check the config file to make sure the right input is being read.")
            kwargs['data']=data
            kwargs['target']=target
            trained_task = self.execute_train_task_with_dependencies(task, **kwargs)
            self.trained_tasks.append(trained_task)
            #If the trained task alters the data in any way, pass it down the chain to the next task
            if hasattr(trained_task, 'data'):
                self.reformatted_input[task.data_format]['data'] = trained_task.data
        log.info("Finished training.")
1192916
def train(X_train, X_test, y_train, y_test, **kwargs):
    '''
    
    '''
    crf = Trainer()
    crf.set_params({
        'c1': 1.0,
        'c2': 0.001,
        'max_iterations': 200,
    })

    for xseq, yseq in zip(X_train, y_train):
        crf.append(xseq, yseq)
    crf.train(model_name)
    return crf
106210
def train(logdir):
  """"""

  ptb_train = PTB(
      name="ptb_train",
      subset="train",
      seq_len=FLAGS.unroll_steps,
      batch_size=FLAGS.batch_size)

  # Connect to training set.
  data_ops = ptb_train()
  embed_layer, rnn_core, output_linear = build_modules(
      is_training=True, vocab_size=ptb_train.vocab_size)
  prediction_logits, zero_state_op = build_logits(
      data_ops, embed_layer, rnn_core, output_linear, name_prefix="train")
  data_loss = build_loss(prediction_logits, data_ops.sparse_target)

  # Add the KL cost.
  total_kl_cost = bbb.get_total_kl_cost()
  num_dataset_elements = FLAGS.batch_size * ptb_train.num_batches
  scaled_kl_cost = total_kl_cost / num_dataset_elements
  total_loss = tf.add(scaled_kl_cost, data_loss)

  # Optimize as usual.
  global_step = tf.get_variable(
      "num_weight_updates",
      initializer=tf.constant(0, dtype=tf.int32, shape=()),
      collections=[tf.GraphKeys.GLOBAL_VARIABLES, tf.GraphKeys.GLOBAL_STEP])

  learning_rate = tf.get_variable(
      "lr", initializer=tf.constant(FLAGS.lr_start, shape=(), dtype=tf.float32))
  learning_rate_update = learning_rate.assign(learning_rate * FLAGS.lr_decay)
  optimizer = tf.train.GradientDescentOptimizer(
      learning_rate=learning_rate)
  optimizer = GlobalNormClippingOptimizer(optimizer, clip_norm=5.0)

  with tf.control_dependencies([optimizer.minimize(total_loss)]):
    global_step_and_train = global_step.assign_add(1)

  # Connect to valid set.
  ptb_valid = PTB(
      name="ptb_valid",
      subset="valid",
      seq_len=FLAGS.unroll_steps,
      batch_size=FLAGS.batch_size)
  valid_data_ops = ptb_valid()
  valid_logits, zero_valid_state = build_logits(
      valid_data_ops, embed_layer, rnn_core, output_linear, name_prefix="valid")
  valid_loss = build_loss(valid_logits, valid_data_ops.sparse_target)

  # Compute metrics for the sake of monitoring training.
  predictions = tf.cast(
      tf.argmax(prediction_logits, axis=-1), tf.int32, name="pred")
  correct_prediction_mask = tf.cast(
      tf.equal(predictions, data_ops.sparse_target), tf.int32)
  accuracy = tf.reduce_mean(
      tf.cast(correct_prediction_mask, tf.float32), name="acc")
  error_rate = tf.subtract(1.0, accuracy, name="err")
  label_probs = tf.nn.softmax(prediction_logits, dim=-1)
  predictive_entropy = tf.reduce_mean(
      label_probs * tf.log(label_probs + 1e-12) * -1.0)

  # Create tf.summary ops.
  log_ops_to_run = {
      "scalar": collections.OrderedDict([
          ("task_loss", data_loss),
          ("train_err_rate", error_rate),
          ("pred_entropy", predictive_entropy),
          ("learning_rate", learning_rate),
          ("elbo_loss", total_loss),
          ("kl_cost", total_kl_cost),
          ("scaled_kl_cost", scaled_kl_cost),
      ]),
      "text": collections.OrderedDict([
          ("labels", ptb_train.to_string_tensor(data_ops.sparse_target)),
          ("predictions", ptb_train.to_string_tensor(predictions))
      ])
  }

  for name, tensor in log_ops_to_run["scalar"].items():
    tf.summary.scalar(os.path.join("train", name), tensor)

  # The remaining logic runs the training loop and logging.
  summary_writer = tf.summary.FileWriterCache.get(logdir=logdir)
  tf.logging.info(
      "Beginning training for {} epochs, each with {} batches.".format(
          FLAGS.num_training_epochs, ptb_train.num_batches))
  with tf.train.MonitoredTrainingSession(
      is_chief=True, checkpoint_dir=logdir, save_summaries_secs=10) as sess:
    num_updates_v = _run_session_with_no_hooks(sess, global_step)
    epoch_idx_start, step_idx_start = divmod(
        num_updates_v, ptb_train.num_batches)
    tf.logging.info("On start, epoch: {}\t step: {}".format(
        epoch_idx_start, step_idx_start))
    for epoch_idx in range(epoch_idx_start, FLAGS.num_training_epochs):
      tf.logging.info("Beginning Epoch {}/{}".format(
          epoch_idx, FLAGS.num_training_epochs))
      tf.logging.info(
          ("Beginning by evaluating on the validation set, which has "
           "{} batches.".format(ptb_valid.num_batches)))
      valid_cost = 0
      valid_steps = 0
      _run_session_with_no_hooks(sess, zero_valid_state)
      for _ in range(ptb_valid.num_batches):
        valid_cost_v, num_updates_v = _run_session_with_no_hooks(
            sess, [valid_loss, global_step])
        valid_cost += valid_cost_v
        valid_steps += FLAGS.unroll_steps
      tf.logging.info("Validation set perplexity: {}".format(
          np.exp(valid_cost / valid_steps)))
      summary = tf.summary.Summary()
      summary.value.add(
          tag="valid/word_level_perplexity",
          simple_value=np.exp(valid_cost / valid_steps))
      summary_writer.add_summary(summary, num_updates_v)

      # Run a training epoch.
      epoch_cost = 0
      epoch_steps = 0
      for batch_idx in range(step_idx_start, ptb_train.num_batches):
        scalars_res, num_updates_v = sess.run(
            [log_ops_to_run["scalar"], global_step_and_train])
        epoch_cost += scalars_res["task_loss"]
        epoch_steps += FLAGS.unroll_steps
        if (batch_idx - 1) % FLAGS.print_every_batches == 0:
          summary = tf.summary.Summary()
          summary.value.add(
              tag="train/word_level_perplexity",
              simple_value=np.exp(epoch_cost / epoch_steps))
          summary_writer.add_summary(summary, num_updates_v)
          scalars_res, strings_res = _run_session_with_no_hooks(
              sess, [log_ops_to_run["scalar"], log_ops_to_run["text"]])
          tf.logging.info("Num weight updates: {}".format(num_updates_v))
          for name, result in scalars_res.items():
            tf.logging.info("{}: {}".format(name, result))
          for name, result in strings_res.items():
            tf.logging.info("{}: {}".format(name, result))

      word_level_perplexity = np.exp(epoch_cost / epoch_steps)
      tf.logging.info(
          "Train Perplexity after Epoch {}: {}".format(
              epoch_idx, word_level_perplexity))

      end_of_epoch_fetches = [zero_state_op]
      if epoch_idx >= FLAGS.high_lr_epochs:
        end_of_epoch_fetches.append(learning_rate_update)
      _run_session_with_no_hooks(sess, end_of_epoch_fetches)

  tf.logging.info("Done training. Thanks for your time.")
81387
def on_train_begin(self, **kwargs):
        ""
        if hasattr(self.loss_func, "reduction") and (self.loss_func.reduction != "sum"):
             warn("For better gradients consider 'reduction=sum'")
321798
def on_batch_begin(self, train, **kwargs):
        ""
        if not self.learn.gan_trainer.gen_mode and train: self.learn.opt.lr *= self.mult_lr
72811
@Override
	public void trainOnInstanceImpl(Instance inst)
	{
		// If this is the first instance, then initialize the forest.
		if(this.numInstances == 0)
		{
			this.buildForest(inst);
		}
		
		// Update the mass profile of every HSTree in the forest
		for(int i = 0 ; i < this.numTrees ; i++)
		{
			forest[i].updateMass(inst, referenceWindow);
		}
		
		if(this.numInstances > 50)
			referenceWindow = false;
		
		// If this is the last instance of the window, update every HSTree's model
		if(this.numInstances % windowSize == 0)
		{
			for(int i = 0 ; i < this.numTrees ; i++)
			{
				forest[i].updateModel();
			}
		}
		
		this.numInstances++;
	}
490406
def train(self):
        
        self.stamp_start = time.time()
        for iteration, batch in tqdm.tqdm(enumerate(self.iter_train),
                                          desc='train', total=self.max_iter,
                                          ncols=80):
            self.epoch = self.iter_train.epoch
            self.iteration = iteration

            ############
            # validate #
            ############

            if self.interval_validate and \
                    self.iteration % self.interval_validate == 0:
                self.validate()

            #########
            # train #
            #########

            batch = map(datasets.transform_lsvrc2012_vgg16, batch)
            in_vars = utils.batch_to_vars(batch, device=self.device)
            self.model.zerograds()
            loss = self.model(*in_vars)

            if loss is not None:
                loss.backward()
                self.optimizer.update()

                lbl_true = zip(*batch)[1]
                lbl_pred = chainer.functions.argmax(self.model.score, axis=1)
                lbl_pred = chainer.cuda.to_cpu(lbl_pred.data)
                acc = utils.label_accuracy_score(
                    lbl_true, lbl_pred, self.model.n_class)
                self._write_log(**{
                    'epoch': self.epoch,
                    'iteration': self.iteration,
                    'elapsed_time': time.time() - self.stamp_start,
                    'train/loss': float(loss.data),
                    'train/acc': acc[0],
                    'train/acc_cls': acc[1],
                    'train/mean_iu': acc[2],
                    'train/fwavacc': acc[3],
                })

            if iteration >= self.max_iter:
                self._save_model()
                break
1330983
def on_train_end(self, **kwargs):
        ""

        if self.save_model:
            # Adapted from fast.ai "SaveModelCallback"
            if self.model_path.is_file():
                with self.model_path.open('rb') as model_file:
                    self.learn.load(model_file, purge=False)
                    print(f'Loaded best saved model from {self.model_path}')
221528
def train(net, X_train, y_train, epochs, verbose_epoch, learning_rate,
          weight_decay, batch_size):
    """"""
    dataset_train = gluon.data.ArrayDataset(X_train, y_train)
    data_iter_train = gluon.data.DataLoader(dataset_train, batch_size,
                                            shuffle=True)
    trainer = gluon.Trainer(net.collect_params(), 'adam',
                            {'learning_rate': learning_rate,
                             'wd': weight_decay})
    net.initialize(force_reinit=True)
    for epoch in range(epochs):
        for data, label in data_iter_train:
            with autograd.record():
                output = net(data)
                loss = square_loss(output, label)
            loss.backward()
            trainer.step(batch_size)
            avg_loss = get_rmse_log(net, X_train, y_train)
        if epoch > verbose_epoch:
            print("Epoch %d, train loss: %f" % (epoch, avg_loss))
    return avg_loss
815009
def train(net, train_data, test_data):
    """"""
    start_pipeline_time = time.time()
    net, trainer = text_cnn.init(net, vocab, args.model_mode, context, args.lr)
    random.shuffle(train_data)
    sp = int(len(train_data)*0.9)
    train_dataloader = DataLoader(dataset=train_data[:sp],
                                  batch_size=args.batch_size,
                                  shuffle=True)
    val_dataloader = DataLoader(dataset=train_data[sp:],
                                batch_size=args.batch_size,
                                shuffle=False)
    test_dataloader = DataLoader(dataset=test_data,
                                 batch_size=args.batch_size,
                                 shuffle=False)
    # Training/Testing.
    best_val_acc = 0
    for epoch in range(args.epochs):
        # Epoch training stats.
        start_epoch_time = time.time()
        epoch_L = 0.0
        epoch_sent_num = 0
        epoch_wc = 0
        # Log interval training stats.
        start_log_interval_time = time.time()
        log_interval_wc = 0
        log_interval_sent_num = 0
        log_interval_L = 0.0
        for i, (data, label) in enumerate(train_dataloader):
            data = mx.nd.transpose(data.as_in_context(context))
            label = label.as_in_context(context)
            wc = max_len
            log_interval_wc += wc
            epoch_wc += wc
            log_interval_sent_num += data.shape[1]
            epoch_sent_num += data.shape[1]

            with autograd.record():
                output = net(data)
                L = loss(output, label).mean()
            L.backward()
            # Update parameter.
            trainer.step(1)
            log_interval_L += L.asscalar()
            epoch_L += L.asscalar()
            if (i + 1) % args.log_interval == 0:
                print('[Epoch %d Batch %d/%d] avg loss %g, throughput %gK wps' % (
                    epoch, i + 1, len(train_dataloader),
                    log_interval_L / log_interval_sent_num,
                    log_interval_wc / 1000 / (time.time() - start_log_interval_time)))
                # Clear log interval training stats.
                start_log_interval_time = time.time()
                log_interval_wc = 0
                log_interval_sent_num = 0
                log_interval_L = 0
        end_epoch_time = time.time()
        val_avg_L, val_acc = evaluate(net, val_dataloader)
        print('[Epoch %d] train avg loss %g, '
              'test acc %.4f, test avg loss %g, throughput %gK wps' % (
                  epoch, epoch_L / epoch_sent_num,
                  val_acc, val_avg_L,
                  epoch_wc / 1000 / (end_epoch_time - start_epoch_time)))

        if val_acc >= best_val_acc:
            print('Observed Improvement.')
            best_val_acc = val_acc
            test_avg_L, test_acc = evaluate(net, test_dataloader)

    print('Test loss %g, test acc %.4f'%(test_avg_L, test_acc))
    print('Total time cost %.2fs'%(time.time()-start_pipeline_time))
    return test_acc
1325949
def train(data_train, model, nsp_loss, mlm_loss, vocab_size, ctx):
    """"""
    hvd.broadcast_parameters(model.collect_params(), root_rank=0)

    mlm_metric = nlp.metric.MaskedAccuracy()
    nsp_metric = nlp.metric.MaskedAccuracy()
    mlm_metric.reset()
    nsp_metric.reset()

    logging.debug('Creating distributed trainer...')
    lr = args.lr
    optim_params = {'learning_rate': lr, 'epsilon': 1e-6, 'wd': 0.01}
    if args.dtype == 'float16':
        optim_params['multi_precision'] = True

    dynamic_loss_scale = args.dtype == 'float16'
    if dynamic_loss_scale:
        loss_scale_param = {'scale_window': 2000 / num_workers}
    else:
        loss_scale_param = None
    trainer = hvd.DistributedTrainer(model.collect_params(), 'bertadam', optim_params)
    fp16_trainer = FP16Trainer(trainer, dynamic_loss_scale=dynamic_loss_scale,
                               loss_scaler_params=loss_scale_param)

    if args.ckpt_dir and args.start_step:
        trainer.load_states(os.path.join(args.ckpt_dir, '%07d.states'%args.start_step))

    accumulate = args.accumulate
    num_train_steps = args.num_steps
    warmup_ratio = args.warmup_ratio
    num_warmup_steps = int(num_train_steps * warmup_ratio)
    params = [p for p in model.collect_params().values() if p.grad_req != 'null']
    param_dict = model.collect_params()

    # Do not apply weight decay on LayerNorm and bias terms
    for _, v in model.collect_params('.*beta|.*gamma|.*bias').items():
        v.wd_mult = 0.0
    if accumulate > 1:
        for p in params:
            p.grad_req = 'add'

    train_begin_time = time.time()
    begin_time = time.time()
    running_mlm_loss, running_nsp_loss = 0, 0
    running_num_tks = 0
    batch_num = 0
    step_num = args.start_step

    logging.debug('Training started')
    while step_num < num_train_steps:
        for _, dataloader in enumerate(data_train):
            if step_num >= num_train_steps:
                break

            # create dummy data loader if needed
            if args.dummy_data_len:
                target_shape = (args.batch_size, args.dummy_data_len)
                dataloader = get_dummy_dataloader(dataloader, target_shape)

            for _, data_batch in enumerate(dataloader):
                if step_num >= num_train_steps:
                    break
                if batch_num % accumulate == 0:
                    step_num += 1
                    # if accumulate > 1, grad_req is set to 'add', and zero_grad is required
                    if accumulate > 1:
                        param_dict.zero_grad()
                    # update learning rate
                    if step_num <= num_warmup_steps:
                        new_lr = lr * step_num / num_warmup_steps
                    else:
                        offset = lr * step_num / num_train_steps
                        new_lr = lr - offset
                    trainer.set_learning_rate(new_lr)
                    if args.profile:
                        profile(step_num, 10, 14, profile_name=args.profile + str(rank))

                # load data
                if args.use_avg_len:
                    data_list = [[seq.as_in_context(context) for seq in shard]
                                 for context, shard in zip([ctx], data_batch)]
                else:
                    data_list = list(split_and_load(data_batch, [ctx]))
                data = data_list[0]

                # forward
                with mx.autograd.record():
                    (ls, ns_label, classified, masked_id, decoded, \
                     masked_weight, ls1, ls2, valid_len) = forward(data, model, mlm_loss,
                                                                   nsp_loss, vocab_size, args.dtype)
                    ls = ls / accumulate
                    # backward
                    if args.dtype == 'float16':
                        fp16_trainer.backward(ls)
                    else:
                        ls.backward()

                running_mlm_loss += ls1.as_in_context(mx.cpu())
                running_nsp_loss += ls2.as_in_context(mx.cpu())
                running_num_tks += valid_len.sum().as_in_context(mx.cpu())

                # update
                if (batch_num + 1) % accumulate == 0:
                    # step() performs 3 things:
                    # 1. allreduce gradients from all workers
                    # 2. checking the global_norm of gradients and clip them if necessary
                    # 3. averaging the gradients and apply updates
                    fp16_trainer.step(1, max_norm=1*num_workers)

                nsp_metric.update([ns_label], [classified])
                mlm_metric.update([masked_id], [decoded], [masked_weight])

                # logging
                if (step_num + 1) % (args.log_interval) == 0 and (batch_num + 1) % accumulate == 0:
                    log(begin_time, running_num_tks, running_mlm_loss / accumulate,
                        running_nsp_loss / accumulate, step_num, mlm_metric, nsp_metric,
                        trainer, args.log_interval)
                    begin_time = time.time()
                    running_mlm_loss = running_nsp_loss = running_num_tks = 0
                    mlm_metric.reset_local()
                    nsp_metric.reset_local()

                # saving checkpoints
                if args.ckpt_dir and (step_num + 1) % (args.ckpt_interval) == 0 \
                   and (batch_num + 1) % accumulate == 0 and local_rank == 0:
                    save_params(step_num, model, trainer, args.ckpt_dir)

                batch_num += 1

    if local_rank == 0:
        save_params(step_num, model, trainer, args.ckpt_dir)
    mx.nd.waitall()
    train_end_time = time.time()
    logging.info('Train cost={:.1f}s'.format(train_end_time - train_begin_time))
1189111
def create(model_config, path, num_workers, batch_size, augmentations=None, tta=None):
    
    if not os.path.isabs(path):
        path = model_config.project_top_dir(path)

    train_path = os.path.join(path, 'train')
    valid_path = os.path.join(path, 'valid')

    train_ds = ImageDirSource(train_path)
    val_ds = ImageDirSource(valid_path)

    return TrainingData(
        train_ds,
        val_ds,
        num_workers=num_workers,
        batch_size=batch_size,
        augmentations=augmentations,
        # test_time_augmentation=tta
    )
1132495
def initialize(self, training_info, model, environment, device):
        
        self.target_model = self.model_factory.instantiate(action_space=environment.action_space).to(device)
        self.target_model.load_state_dict(model.state_dict())
        self.target_model.eval()
378278
def train(self, sentences, model):
        
        trainer = pycrfsuite.Trainer(verbose=True)
        trainer.set_params(self.params)
        for sentence in sentences:
            tokens, labels = zip(*sentence)
            features = [self._get_features(tokens, i) for i in range(len(tokens))]
            trainer.append(features, labels)
        trainer.train(model)
        self.load(model)
682089
def train(target, dataset, cluster_spec, ctx):
  """"""
  # Number of workers and parameter servers are infered from the workers and ps
  # hosts string.
  num_workers = len(cluster_spec.as_dict()['worker'])
  num_parameter_servers = len(cluster_spec.as_dict()['ps'])
  # If no value is given, num_replicas_to_aggregate defaults to be the number of
  # workers.
  if FLAGS.num_replicas_to_aggregate == -1:
    num_replicas_to_aggregate = num_workers
  else:
    num_replicas_to_aggregate = FLAGS.num_replicas_to_aggregate

  # Both should be greater than 0 in a distributed training.
  assert num_workers > 0 and num_parameter_servers > 0, (' num_workers and '
                                                         'num_parameter_servers'
                                                         ' must be > 0.')

  # Choose worker 0 as the chief. Note that any worker could be the chief
  # but there should be only one chief.
  is_chief = (FLAGS.task_id == 0)

  # Ops are assigned to worker by default.
  with tf.device('/job:worker/task:%d' % FLAGS.task_id):
    # Variables and its related init/assign ops are assigned to ps.
    with slim.scopes.arg_scope(
          [slim.variables.variable, slim.variables.global_step],
          device=slim.variables.VariableDeviceChooser(num_parameter_servers)):
      # Create a variable to count the number of train() calls. This equals the
      # number of updates applied to the variables.
      global_step = slim.variables.global_step()

      # Calculate the learning rate schedule.
      num_batches_per_epoch = (dataset.num_examples_per_epoch() /
                               FLAGS.batch_size)
      # Decay steps need to be divided by the number of replicas to aggregate.
      decay_steps = int(num_batches_per_epoch * FLAGS.num_epochs_per_decay /
                        num_replicas_to_aggregate)

      # Decay the learning rate exponentially based on the number of steps.
      lr = tf.train.exponential_decay(FLAGS.initial_learning_rate,
                                      global_step,
                                      decay_steps,
                                      FLAGS.learning_rate_decay_factor,
                                      staircase=True)
      # Add a summary to track the learning rate.
      tf.summary.scalar('learning_rate', lr)

      # Create an optimizer that performs gradient descent.
      opt = tf.train.RMSPropOptimizer(lr,
                                      RMSPROP_DECAY,
                                      momentum=RMSPROP_MOMENTUM,
                                      epsilon=RMSPROP_EPSILON)

      if FLAGS.input_mode == 'spark':
        def feed_dict(feed_batch):
          # extract TFRecords, since feed_batch is [(TFRecord, None)]
          tfrecords = []
          for elem in feed_batch:
            tfrecords.append(str(elem[0]))
          return tfrecords

        batch = tf.placeholder(tf.string, [FLAGS.batch_size / FLAGS.num_preprocess_threads])

        # The following is adapted from image_processing.py to remove Readers/QueueRunners.
        # Note: this removes the RandomShuffledQueue, so the incoming data is not shuffled.
        # Presumably, this could be done on the Spark side or done in additional TF code.
        examples = tf.unstack(batch)
        images, labels = [], []
        for example_serialized in examples:
          for thread_id in range(FLAGS.num_preprocess_threads):
            # Parse a serialized Example proto to extract the image and metadata.
            image_buffer, label_index, bbox, _ = image_processing.parse_example_proto(example_serialized)
            image = image_processing.image_preprocessing(image_buffer, bbox, train, thread_id)
            images.append(image)
            labels.append(label_index)
        height = FLAGS.image_size
        width = FLAGS.image_size
        depth = 3
        images = tf.cast(images, tf.float32)
        images = tf.reshape(images, shape=[FLAGS.batch_size, height, width, depth])
        tf.summary.image('images', images)
        labels = tf.reshape(labels, [FLAGS.batch_size])
      else:
        images, labels = image_processing.distorted_inputs(
            dataset,
            batch_size=FLAGS.batch_size,
            num_preprocess_threads=FLAGS.num_preprocess_threads)

      # Number of classes in the Dataset label set plus 1.
      # Label 0 is reserved for an (unused) background class.
      num_classes = dataset.num_classes() + 1
      logits = inception.inference(images, num_classes, for_training=True)
      # Add classification loss.
      inception.loss(logits, labels)

      # Gather all of the losses including regularization losses.
      losses = tf.get_collection(slim.losses.LOSSES_COLLECTION)
      losses += tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)

      total_loss = tf.add_n(losses, name='total_loss')

      if is_chief:
        # Compute the moving average of all individual losses and the
        # total loss.
        loss_averages = tf.train.ExponentialMovingAverage(0.9, name='avg')
        loss_averages_op = loss_averages.apply(losses + [total_loss])

        # Attach a scalar summmary to all individual losses and the total loss;
        # do the same for the averaged version of the losses.
        for l in losses + [total_loss]:
          loss_name = l.op.name
          # Name each loss as '(raw)' and name the moving average version of the
          # loss as the original loss name.
          tf.summary.scalar(loss_name + ' (raw)', l)
          tf.summary.scalar(loss_name, loss_averages.average(l))

        # Add dependency to compute loss_averages.
        with tf.control_dependencies([loss_averages_op]):
          total_loss = tf.identity(total_loss)

      # Track the moving averages of all trainable variables.
      # Note that we maintain a 'double-average' of the BatchNormalization
      # global statistics.
      # This is not needed when the number of replicas are small but important
      # for synchronous distributed training with tens of workers/replicas.
      exp_moving_averager = tf.train.ExponentialMovingAverage(
          inception.MOVING_AVERAGE_DECAY, global_step)

      variables_to_average = (
          tf.trainable_variables() + tf.moving_average_variables())

      # Add histograms for model variables.
      for var in variables_to_average:
        tf.summary.histogram(var.op.name, var)

      # Create synchronous replica optimizer.
      opt = tf.train.SyncReplicasOptimizer(
          opt,
          replicas_to_aggregate=num_replicas_to_aggregate,
          total_num_replicas=num_workers,
          variable_averages=exp_moving_averager,
          variables_to_average=variables_to_average)

      batchnorm_updates = tf.get_collection(slim.ops.UPDATE_OPS_COLLECTION)
      assert batchnorm_updates, 'Batchnorm updates are missing'
      batchnorm_updates_op = tf.group(*batchnorm_updates)
      # Add dependency to compute batchnorm_updates.
      with tf.control_dependencies([batchnorm_updates_op]):
        total_loss = tf.identity(total_loss)

      # Compute gradients with respect to the loss.
      grads = opt.compute_gradients(total_loss)

      # Add histograms for gradients.
      for grad, var in grads:
        if grad is not None:
          tf.summary.histogram(var.op.name + '/gradients', grad)

      apply_gradients_op = opt.apply_gradients(grads, global_step=global_step)

      with tf.control_dependencies([apply_gradients_op]):
        train_op = tf.identity(total_loss, name='train_op')

      # Get chief queue_runners, init_tokens and clean_up_op, which is used to
      # synchronize replicas.
      # More details can be found in sync_replicas_optimizer.
      chief_queue_runners = [opt.get_chief_queue_runner()]
      init_tokens_op = opt.get_init_tokens_op()

      # Create a saver.
      saver = tf.train.Saver()

      # Build the summary operation based on the TF collection of Summaries.
      summary_op = tf.summary.merge_all()

      # Build an initialization operation to run below.
      init_op = tf.global_variables_initializer()

      # We run the summaries in the same thread as the training operations by
      # passing in None for summary_op to avoid a summary_thread being started.
      # Running summaries and training operations in parallel could run out of
      # GPU memory.
      summary_writer = tf.summary.FileWriter("tensorboard_%d" % ctx.worker_num, graph=tf.get_default_graph())
      sv = tf.train.Supervisor(is_chief=is_chief,
                               logdir=FLAGS.train_dir,
                               init_op=init_op,
                               summary_op=None,
                               global_step=global_step,
                               summary_writer=summary_writer,
                               saver=saver,
                               save_model_secs=FLAGS.save_interval_secs)

      tf.logging.info('%s Supervisor' % datetime.now())

      sess_config = tf.ConfigProto(
          allow_soft_placement=True,
          log_device_placement=FLAGS.log_device_placement)

      # Get a session.
      sess = sv.prepare_or_wait_for_session(target, config=sess_config)

      # Start the queue runners.
      queue_runners = tf.get_collection(tf.GraphKeys.QUEUE_RUNNERS)
      sv.start_queue_runners(sess, queue_runners)
      tf.logging.info('Started %d queues for processing input data.',
                      len(queue_runners))

      if is_chief:
        sv.start_queue_runners(sess, chief_queue_runners)
        sess.run(init_tokens_op)

      # Train, checking for Nans. Concurrently run the summary operation at a
      # specified interval. Note that the summary_op and train_op never run
      # simultaneously in order to prevent running out of GPU memory.
      next_summary_time = time.time() + FLAGS.save_summaries_secs
      tf_feed = TFNode.DataFeed(ctx.mgr)
      while not sv.should_stop():
        try:
          start_time = time.time()
          if FLAGS.input_mode == 'spark':
            tmp = feed_dict(tf_feed.next_batch(FLAGS.batch_size / FLAGS.num_preprocess_threads))
            feed = {batch: tmp}
            loss_value, step = sess.run([train_op, global_step], feed_dict=feed)
          else:
            loss_value, step = sess.run([train_op, global_step])
          assert not np.isnan(loss_value), 'Model diverged with loss = NaN'
          if step > FLAGS.max_steps:
            break
          duration = time.time() - start_time

          if step % 30 == 0:
            examples_per_sec = FLAGS.batch_size / float(duration)
            format_str = ('Worker %d: %s: step %d, loss = %.2f'
                          '(%.1f examples/sec; %.3f  sec/batch)')
            tf.logging.info(format_str %
                            (FLAGS.task_id, datetime.now(), step, loss_value,
                             examples_per_sec, duration))

          # Determine if the summary_op should be run on the chief worker.
          if FLAGS.input_mode == 'tf' and is_chief and next_summary_time < time.time():
            tf.logging.info('Running Summary operation on the chief.')
            summary_str = sess.run(summary_op)
            sv.summary_computed(sess, summary_str)
            tf.logging.info('Finished running Summary operation.')

            # Determine the next time for running the summary.
            next_summary_time += FLAGS.save_summaries_secs
        except:
          if is_chief:
            tf.logging.info('About to execute sync_clean_up_op!')
          raise

      # Stop the TFNode data feed
      if FLAGS.input_mode == 'spark':
        tf_feed.terminate()

      # Stop the supervisor.  This also waits for service threads to finish.
      sv.stop()

      # Save after the training ends.
      if is_chief:
        saver.save(sess,
                   os.path.join(FLAGS.train_dir, 'model.ckpt'),
                   global_step=global_step)
931036
def train(self, params, epoch):
    

    # Callback used to log information on every batch
    def log(model, batch_idx):
      if batch_idx % params["log_interval"] == 0:
        entropy = model.entropy()
        print("logging: {} learning iterations, entropy: {} / {}".format(
          model.getLearningIterations(), float(entropy), model.maxEntropy()))

        if params["create_plots"]:
          plotDutyCycles(model.dutyCycle,
                         self.resultsDir + "/figure_" + str(epoch) + "_" +
                         str(model.getLearningIterations()))


    # Adjust first epoch batch size to stabilize the dutycycles at the
    # beginning of the training
    loader = self.train_loader
    batches_in_epoch = params["batches_in_epoch"]
    if "first_epoch_batch_size" in params:
      if epoch == 0:
        batches_in_epoch = params.get("batches_in_first_epoch", batches_in_epoch)
        loader = torch.utils.data.DataLoader(self.train_loader.dataset,
                                             batch_size=params["first_epoch_batch_size"],
                                             sampler=self.train_loader.sampler)

    trainModel(model=self.model, loader=loader,
               optimizer=self.optimizer, device=self.device,
               batches_in_epoch=batches_in_epoch,
               batch_callback=log)

    self.model.postEpoch()