Put `learn` on distributed training with `cuda_id`.
348562
def _learner_distributed(learn:Learner, cuda_id:int, cache_dir:PathOrStr='tmp'):
    ""
    learn.callbacks.append(DistributedTrainer(learn, cuda_id))
    learn.callbacks.append(DistributedRecorder(learn, cuda_id, cache_dir))
    return learn
967924
def select_device_by_load(wproc=0.5, wmem=0.5):
    

    ids = device_by_load(wproc=wproc, wmem=wmem)
    cp.cuda.Device(ids[0]).use()
    return ids[0]
1180998
def train(args):
    """"""
    vocab, row, col, counts = get_train_data(args)
    model = GloVe(token_to_idx=vocab.token_to_idx, output_dim=args.emsize,
                  dropout=args.dropout, x_max=args.x_max, alpha=args.alpha,
                  weight_initializer=mx.init.Uniform(scale=1 / args.emsize))
    context = get_context(args)
    model.initialize(ctx=context)
    if not args.no_hybridize:
        model.hybridize(static_alloc=not args.no_static_alloc)

    optimizer_kwargs = dict(learning_rate=args.lr, eps=args.adagrad_eps)
    params = list(model.collect_params().values())
    try:
        trainer = mx.gluon.Trainer(params, 'groupadagrad', optimizer_kwargs)
    except ValueError:
        logging.warning('MXNet <= v1.3 does not contain '
                        'GroupAdaGrad support. Falling back to AdaGrad')
        trainer = mx.gluon.Trainer(params, 'adagrad', optimizer_kwargs)

    index_dtype = 'int32'
    if counts.shape[0] >= np.iinfo(np.int32).max:
        index_dtype = 'int64'
        logging.info('Co-occurrence matrix is large. '
                     'Using int64 to represent sample indices.')
    indices = mx.nd.arange(counts.shape[0], dtype=index_dtype)
    for epoch in range(args.epochs):
        # Logging variables
        log_wc = 0
        log_start_time = time.time()
        log_avg_loss = 0

        mx.nd.shuffle(indices, indices)  # inplace shuffle
        bs = args.batch_size
        num_batches = indices.shape[0] // bs
        for i in range(num_batches):
            batch_indices = indices[bs * i:bs * (i + 1)]
            ctx = context[i % len(context)]
            batch_row = row[batch_indices].as_in_context(ctx)
            batch_col = col[batch_indices].as_in_context(ctx)
            batch_counts = counts[batch_indices].as_in_context(ctx)
            with mx.autograd.record():
                loss = model(batch_row, batch_col, batch_counts)
                loss.backward()

            if len(context) == 1 or (i + 1) % len(context) == 0:
                trainer.step(batch_size=1)

            # Logging
            log_wc += loss.shape[0]
            log_avg_loss += loss.mean().as_in_context(context[0])
            if (i + 1) % args.log_interval == 0:
                # Forces waiting for computation by computing loss value
                log_avg_loss = log_avg_loss.asscalar() / args.log_interval
                wps = log_wc / (time.time() - log_start_time)
                logging.info('[Epoch {} Batch {}/{}] loss={:.4f}, '
                             'throughput={:.2f}K wps, wc={:.2f}K'.format(
                                 epoch, i + 1, num_batches, log_avg_loss,
                                 wps / 1000, log_wc / 1000))
                log_dict = dict(
                    global_step=epoch * len(indices) + i * args.batch_size,
                    epoch=epoch, batch=i + 1, loss=log_avg_loss,
                    wps=wps / 1000)
                log(args, log_dict)

                log_start_time = time.time()
                log_avg_loss = 0
                log_wc = 0

            if args.eval_interval and (i + 1) % args.eval_interval == 0:
                with print_time('mx.nd.waitall()'):
                    mx.nd.waitall()
                with print_time('evaluate'):
                    evaluate(args, model, vocab, i + num_batches * epoch)

    # Evaluate
    with print_time('mx.nd.waitall()'):
        mx.nd.waitall()
    with print_time('evaluate'):
        evaluate(args, model, vocab, num_batches * args.epochs,
                 eval_analogy=not args.no_eval_analogy)

    # Save params
    with print_time('save parameters'):
        model.save_parameters(os.path.join(args.logdir, 'glove.params'))
572675
def cudnnSetStream(handle, id):
    

    status = _libcudnn.cudnnSetStream(handle, id)
    cudnnCheckStatus(status)
230173
def single_node_env(num_gpus=1):
  """"""
  import tensorflow as tf
  # ensure expanded CLASSPATH w/o glob characters (required for Spark 2.1 + JNI)
  if 'HADOOP_PREFIX' in os.environ and 'TFOS_CLASSPATH_UPDATED' not in os.environ:
      classpath = os.environ['CLASSPATH']
      hadoop_path = os.path.join(os.environ['HADOOP_PREFIX'], 'bin', 'hadoop')
      hadoop_classpath = subprocess.check_output([hadoop_path, 'classpath', '--glob']).decode()
      os.environ['CLASSPATH'] = classpath + os.pathsep + hadoop_classpath
      os.environ['TFOS_CLASSPATH_UPDATED'] = '1'

  # reserve GPU, if requested
  if tf.test.is_built_with_cuda():
    gpus_to_use = gpu_info.get_gpus(num_gpus)
    logging.info("Using gpu(s): {0}".format(gpus_to_use))
    os.environ['CUDA_VISIBLE_DEVICES'] = gpus_to_use
  else:
    # CPU
    logging.info("Using CPU")
    os.environ['CUDA_VISIBLE_DEVICES'] = ''
965436
def get_training_data(batch_size):
    
    return gluon.data.DataLoader(
        CIFAR10(train=True, transform=transformer),
        batch_size=batch_size, shuffle=True, last_batch='discard')
1360164
def load(self, filename):
        
        if os.path.isfile(filename):
            checkpoint = torch.load(filename, map_location={'cuda:0': 'cpu'})
            if self.distributed:
                self.model.module.load_state_dict(checkpoint['state_dict'])
            else:
                self.model.load_state_dict(checkpoint['state_dict'])
            self.fp_optimizer.initialize_model(self.model)
            self.optimizer.load_state_dict(checkpoint['optimizer'])
            self.scheduler.load_state_dict(checkpoint['scheduler'])
            self.epoch = checkpoint['epoch']
            self.loss = checkpoint['loss']
            logging.info(f'Loaded checkpoint {filename} (epoch {self.epoch})')
        else:
            logging.error(f'Invalid checkpoint: {filename}')
480479
def init_train_args(self, parser):
        ''' '''
        parser.add_argument('-f', '--configuration', dest='config', default=DEFAULT_USER_CONFIG_PATH,
                            help='the path to the configuration file to use -- ./config.yaml by default')
        parser.add_argument('-c', '--dataset-name', help='the name of the saved dataset to train on')
        parser.add_argument('-s', '--saved-model-name', help='the name of the saved model to continue training')
        parser.add_argument('-n', '--new-model-name', help='the name to give the newly trained model')
        parser.add_argument('-g', '--gpus', dest='num_gpus', type=int, default=1, help='the number of gpus to use during training')
981932
def train(num_training_iterations, report_interval,
          reduce_learning_rate_interval):
  """"""

  # Build the computation graph.
  graph_tensors, dataset_train = build_graph(
      lstm_depth=FLAGS.lstm_depth, batch_size=FLAGS.batch_size,
      num_embedding=FLAGS.num_embedding, num_hidden=FLAGS.num_hidden,
      truncation_length=FLAGS.truncation_length,
      sample_length=FLAGS.sample_length, max_grad_norm=FLAGS.max_grad_norm,
      initial_learning_rate=FLAGS.learning_rate,
      reduce_learning_rate_multiplier=FLAGS.reduce_learning_rate_multiplier,
      optimizer_epsilon=FLAGS.optimizer_epsilon)

  # Configure a checkpoint saver.
  saver_hook = _configure_saver(FLAGS.checkpoint_dir,
                                FLAGS.checkpoint_interval)

  # Train the network.
  with tf.train.SingularMonitoredSession(
      hooks=[saver_hook], checkpoint_dir=FLAGS.checkpoint_dir) as sess:

    start_iteration = sess.run(graph_tensors["global_step"])

    for train_iteration in range(start_iteration, num_training_iterations):
      if (train_iteration + 1) % report_interval == 0:
        train_loss_v, valid_loss_v, _ = sess.run(
            (graph_tensors["train_loss"],
             graph_tensors["valid_loss"],
             graph_tensors["train_step"]))

        train_generated_string_v = sess.run(
            graph_tensors["train_generated_string"])

        train_generated_string_human = dataset_train.to_human_readable(
            (train_generated_string_v, 0), indices=[0])

        tf.logging.info("%d: Training loss %f. Validation loss %f. Sample = %s",
                        train_iteration,
                        train_loss_v,
                        valid_loss_v,
                        train_generated_string_human)
      else:
        train_loss_v, _ = sess.run((graph_tensors["train_loss"],
                                    graph_tensors["train_step"]))
        tf.logging.info("%d: Training loss %f.", train_iteration, train_loss_v)

      if (train_iteration + 1) % reduce_learning_rate_interval == 0:
        sess.run(graph_tensors["reduce_learning_rate"])
        tf.logging.info("Reducing learning rate.")

    test_loss = sess.run(graph_tensors["test_loss"])
    tf.logging.info("Test loss %f", test_loss)
801268
def main(
        gpu:Param("GPU to run on", str)=None,
        woof: Param("Use imagewoof (otherwise imagenette)", int)=0,
        lr: Param("Learning rate", float)=1e-3,
        size: Param("Size (px: 128,192,224)", int)=128,
        alpha: Param("Alpha", float)=0.99,
        mom: Param("Momentum", float)=0.9,
        eps: Param("epsilon", float)=1e-6,
        epochs: Param("Number of epochs", int)=5,
        bs: Param("Batch size", int)=256,
        mixup: Param("Mixup", float)=0.,
        opt: Param("Optimizer (adam,rms,sgd)", str)='adam',
        arch: Param("Architecture (xresnet34, xresnet50, presnet34, presnet50)", str)='xresnet50',
        dump: Param("Print model; don't train", int)=0,
        ):
    ""

    gpu = setup_distrib(gpu)
    if gpu is None: bs *= torch.cuda.device_count()
    if   opt=='adam' : opt_func = partial(optim.Adam, betas=(mom,alpha), eps=eps)
    elif opt=='rms'  : opt_func = partial(optim.RMSprop, alpha=alpha, eps=eps)
    elif opt=='sgd'  : opt_func = partial(optim.SGD, momentum=mom)

    data = get_data(size, woof, bs)
    bs_rat = bs/256
    if gpu is not None: bs_rat *= num_distrib()
    if not gpu: print(f'lr: {lr}; eff_lr: {lr*bs_rat}; size: {size}; alpha: {alpha}; mom: {mom}; eps: {eps}')
    lr *= bs_rat

    m = globals()[arch]
    learn = (Learner(data, m(c_out=10), wd=1e-2, opt_func=opt_func,
             metrics=[accuracy,top_k_accuracy],
             bn_wd=False, true_wd=True,
             loss_func = LabelSmoothingCrossEntropy())
            )
    if dump: print(learn.model); exit()
    if mixup: learn = learn.mixup(alpha=mixup)
    learn = learn.to_fp16(dynamic=True)
    if gpu is None:       learn.to_parallel()
    elif num_distrib()>1: learn.to_distributed(gpu) # Requires `-m fastai.launch`

    learn.fit_one_cycle(epochs, lr, div_factor=10, pct_start=0.3)
476147
def optimize(self, loss, num_async_replicas=1, use_tpu=False):
    """"""
    lr = learning_rate.learning_rate_schedule(self.hparams)
    if num_async_replicas > 1:
      log_info("Dividing learning rate by num_async_replicas: %d",
               num_async_replicas)
    lr /= math.sqrt(float(num_async_replicas))
    train_op = optimize.optimize(loss, lr, self.hparams, use_tpu=use_tpu)
    return train_op
216663
def build_trainer(nn, ds, verbosity=1):
    """"""
    return pb.supervised.trainers.rprop.RPropMinusTrainer(nn, dataset=ds, batchlearning=True, verbose=bool(verbosity))
388745
def train():
    """"""
    trainer = gluon.Trainer(model.collect_params(), args.optimizer,
                            {'learning_rate': args.lr, 'beta2': 0.98, 'epsilon': 1e-9})

    train_data_loader, val_data_loader, test_data_loader \
        = dataprocessor.make_dataloader(data_train, data_val, data_test, args,
                                        use_average_length=True, num_shards=len(ctx))

    if args.bleu == 'tweaked':
        bpe = bool(args.dataset != 'IWSLT2015' and args.dataset != 'TOY')
        split_compound_word = bpe
        tokenized = True
    elif args.bleu == '13a' or args.bleu == 'intl':
        bpe = False
        split_compound_word = False
        tokenized = False
    else:
        raise NotImplementedError

    best_valid_bleu = 0.0
    step_num = 0
    warmup_steps = args.warmup_steps
    grad_interval = args.num_accumulated
    model.collect_params().setattr('grad_req', 'add')
    average_start = (len(train_data_loader) // grad_interval) * (args.epochs - args.average_start)
    average_param_dict = None
    model.collect_params().zero_grad()
    parallel = Parallel(num_ctxs, parallel_model)
    for epoch_id in range(args.epochs):
        log_avg_loss = 0
        log_wc = 0
        loss_denom = 0
        step_loss = 0
        log_start_time = time.time()
        for batch_id, seqs \
                in enumerate(train_data_loader):
            if batch_id % grad_interval == 0:
                step_num += 1
                new_lr = args.lr / math.sqrt(args.num_units) \
                         * min(1. / math.sqrt(step_num), step_num * warmup_steps ** (-1.5))
                trainer.set_learning_rate(new_lr)
            src_wc, tgt_wc, bs = np.sum([(shard[2].sum(), shard[3].sum(), shard[0].shape[0])
                                         for shard in seqs], axis=0)
            seqs = [[seq.as_in_context(context) for seq in shard]
                    for context, shard in zip(ctx, seqs)]
            Ls = []
            for seq in seqs:
                parallel.put((seq, args.batch_size))
            Ls = [parallel.get() for _ in range(len(ctx))]
            src_wc = src_wc.asscalar()
            tgt_wc = tgt_wc.asscalar()
            loss_denom += tgt_wc - bs
            if batch_id % grad_interval == grad_interval - 1 or\
                    batch_id == len(train_data_loader) - 1:
                if average_param_dict is None:
                    average_param_dict = {k: v.data(ctx[0]).copy() for k, v in
                                          model.collect_params().items()}
                trainer.step(float(loss_denom) / args.batch_size / 100.0)
                param_dict = model.collect_params()
                param_dict.zero_grad()
                if step_num > average_start:
                    alpha = 1. / max(1, step_num - average_start)
                    for name, average_param in average_param_dict.items():
                        average_param[:] += alpha * (param_dict[name].data(ctx[0]) - average_param)
            step_loss += sum([L.asscalar() for L in Ls])
            if batch_id % grad_interval == grad_interval - 1 or\
                    batch_id == len(train_data_loader) - 1:
                log_avg_loss += step_loss / loss_denom * args.batch_size * 100.0
                loss_denom = 0
                step_loss = 0
            log_wc += src_wc + tgt_wc
            if (batch_id + 1) % (args.log_interval * grad_interval) == 0:
                wps = log_wc / (time.time() - log_start_time)
                logging.info('[Epoch {} Batch {}/{}] loss={:.4f}, ppl={:.4f}, '
                             'throughput={:.2f}K wps, wc={:.2f}K'
                             .format(epoch_id, batch_id + 1, len(train_data_loader),
                                     log_avg_loss / args.log_interval,
                                     np.exp(log_avg_loss / args.log_interval),
                                     wps / 1000, log_wc / 1000))
                log_start_time = time.time()
                log_avg_loss = 0
                log_wc = 0
        mx.nd.waitall()
        valid_loss, valid_translation_out = evaluate(val_data_loader, ctx[0])
        valid_bleu_score, _, _, _, _ = compute_bleu([val_tgt_sentences], valid_translation_out,
                                                    tokenized=tokenized, tokenizer=args.bleu,
                                                    split_compound_word=split_compound_word,
                                                    bpe=bpe)
        logging.info('[Epoch {}] valid Loss={:.4f}, valid ppl={:.4f}, valid bleu={:.2f}'
                     .format(epoch_id, valid_loss, np.exp(valid_loss), valid_bleu_score * 100))
        test_loss, test_translation_out = evaluate(test_data_loader, ctx[0])
        test_bleu_score, _, _, _, _ = compute_bleu([test_tgt_sentences], test_translation_out,
                                                   tokenized=tokenized, tokenizer=args.bleu,
                                                   split_compound_word=split_compound_word,
                                                   bpe=bpe)
        logging.info('[Epoch {}] test Loss={:.4f}, test ppl={:.4f}, test bleu={:.2f}'
                     .format(epoch_id, test_loss, np.exp(test_loss), test_bleu_score * 100))
        dataprocessor.write_sentences(valid_translation_out,
                                      os.path.join(args.save_dir,
                                                   'epoch{:d}_valid_out.txt').format(epoch_id))
        dataprocessor.write_sentences(test_translation_out,
                                      os.path.join(args.save_dir,
                                                   'epoch{:d}_test_out.txt').format(epoch_id))
        if valid_bleu_score > best_valid_bleu:
            best_valid_bleu = valid_bleu_score
            save_path = os.path.join(args.save_dir, 'valid_best.params')
            logging.info('Save best parameters to {}'.format(save_path))
            model.save_parameters(save_path)
        save_path = os.path.join(args.save_dir, 'epoch{:d}.params'.format(epoch_id))
        model.save_parameters(save_path)
    save_path = os.path.join(args.save_dir, 'average.params')
    mx.nd.save(save_path, average_param_dict)
    if args.average_checkpoint:
        for j in range(args.num_averages):
            params = mx.nd.load(os.path.join(args.save_dir,
                                             'epoch{:d}.params'.format(args.epochs - j - 1)))
            alpha = 1. / (j + 1)
            for k, v in model._collect_params_with_prefix().items():
                for c in ctx:
                    v.data(c)[:] += alpha * (params[k].as_in_context(c) - v.data(c))
        save_path = os.path.join(args.save_dir,
                                 'average_checkpoint_{}.params'.format(args.num_averages))
        model.save_parameters(save_path)
    elif args.average_start > 0:
        for k, v in model.collect_params().items():
            v.set_data(average_param_dict[k])
        save_path = os.path.join(args.save_dir, 'average.params')
        model.save_parameters(save_path)
    else:
        model.load_parameters(os.path.join(args.save_dir, 'valid_best.params'), ctx)
    valid_loss, valid_translation_out = evaluate(val_data_loader, ctx[0])
    valid_bleu_score, _, _, _, _ = compute_bleu([val_tgt_sentences], valid_translation_out,
                                                tokenized=tokenized, tokenizer=args.bleu, bpe=bpe,
                                                split_compound_word=split_compound_word)
    logging.info('Best model valid Loss={:.4f}, valid ppl={:.4f}, valid bleu={:.2f}'
                 .format(valid_loss, np.exp(valid_loss), valid_bleu_score * 100))
    test_loss, test_translation_out = evaluate(test_data_loader, ctx[0])
    test_bleu_score, _, _, _, _ = compute_bleu([test_tgt_sentences], test_translation_out,
                                               tokenized=tokenized, tokenizer=args.bleu, bpe=bpe,
                                               split_compound_word=split_compound_word)
    logging.info('Best model test Loss={:.4f}, test ppl={:.4f}, test bleu={:.2f}'
                 .format(test_loss, np.exp(test_loss), test_bleu_score * 100))
    dataprocessor.write_sentences(valid_translation_out,
                                  os.path.join(args.save_dir, 'best_valid_out.txt'))
    dataprocessor.write_sentences(test_translation_out,
                                  os.path.join(args.save_dir, 'best_test_out.txt'))
328946
def NeuralGPU(feature_depth=96, steps=16, vocab_size=2):
  
  xs = []
  xs.append(
      layers.Embedding(feature_depth=feature_depth, vocab_size=vocab_size))
  core = ConvDiagonalGRU(units=feature_depth)
  xs.extend([core] * steps)
  xs.append(layers.Dense(vocab_size))
  xs.append(layers.LogSoftmax())

  return layers.Serial(*xs)
890524
def run_toy_HMC(gpu_id=None):
    """"""
    X, Y, X_test, Y_test = load_toy()
    minibatch_size = Y.shape[0]
    noise_precision = 1 / 9.0
    net = get_toy_sym(True, noise_precision)
    data_shape = (minibatch_size,) + X.shape[1::]
    data_inputs = {'data': nd.zeros(data_shape, ctx=dev(gpu_id)),
                   'teacher_output_label': nd.zeros((minibatch_size, 1), ctx=dev(gpu_id))}
    initializer = mx.init.Uniform(0.07)
    sample_pool = HMC(net, data_inputs=data_inputs, X=X, Y=Y, X_test=X_test, Y_test=Y_test,
                      sample_num=300000, initializer=initializer, prior_precision=1.0,
                      learning_rate=1E-3, L=10, dev=dev(gpu_id))
655331
def _build_train_op(self):
        """"""
        num_gpus = self.hps.num_gpus if self.hps.num_gpus != 0 else 1
        # The learning rate schedule is dependent on the number of gpus.
        boundaries = [int(20000 * i / np.sqrt(num_gpus)) for i in range(2, 5)]
        values = [0.1, 0.01, 0.001, 0.0001]
        self.lrn_rate = tf.train.piecewise_constant(self.global_step,
                                                    boundaries, values)
        tf.summary.scalar("learning rate", self.lrn_rate)

        if self.hps.optimizer == "sgd":
            optimizer = tf.train.GradientDescentOptimizer(self.lrn_rate)
        elif self.hps.optimizer == "mom":
            optimizer = tf.train.MomentumOptimizer(self.lrn_rate, 0.9)

        apply_op = optimizer.minimize(self.cost, global_step=self.global_step)
        train_ops = [apply_op] + self._extra_train_ops
        self.train_op = tf.group(*train_ops)
        self.variables = ray.experimental.tf_utils.TensorFlowVariables(
            self.train_op)
1209978
def _collect_gpu_info(run_info):
  """"""
  gpu_info = {}
  local_device_protos = device_lib.list_local_devices()

  gpu_info["count"] = len([d for d in local_device_protos
                           if d.device_type == "GPU"])
  # The device description usually is a JSON string, which contains the GPU
  # model info, eg:
  # "device: 0, name: Tesla P100-PCIE-16GB, pci bus id: 0000:00:04.0"
  for d in local_device_protos:
    if d.device_type == "GPU":
      gpu_info["model"] = _parse_gpu_model(d.physical_device_desc)
      # Assume all the GPU connected are same model
      break
  run_info["machine_config"]["gpu_info"] = gpu_info