Build convnet style learner.
484263
def cnn_learner(data:DataBunch, base_arch:Callable, cut:Union[int,Callable]=None, pretrained:bool=True,
                lin_ftrs:Optional[Collection[int]]=None, ps:Floats=0.5, custom_head:Optional[nn.Module]=None,
                split_on:Optional[SplitFuncOrIdxList]=None, bn_final:bool=False, init=nn.init.kaiming_normal_,
                concat_pool:bool=True, **kwargs:Any)->Learner:
    ""
    meta = cnn_config(base_arch)
    model = create_cnn_model(base_arch, data.c, cut, pretrained, lin_ftrs, ps=ps, custom_head=custom_head,
        split_on=split_on, bn_final=bn_final, concat_pool=concat_pool)
    learn = Learner(data, model, **kwargs)
    learn.split(split_on or meta['split'])
    if pretrained: learn.freeze()
    if init: apply_init(model[1], init)
    return learn
1156057
def build_network(self):
        '''
        
        '''

        # Reshape to use within a convolutional neural net.
        # Last dimension is for "features" - there is only one here, since images are
        # grayscale -- it would be 3 for an RGB image, 4 for RGBA, etc.
        with tf.name_scope('reshape'):
            try:
                input_dim = int(math.sqrt(self.x_dim))
            except:
                print(
                    'input dim cannot be sqrt and reshape. input dim: ' + str(self.x_dim))
                logger.debug(
                    'input dim cannot be sqrt and reshape. input dim: %s', str(self.x_dim))
                raise
            x_image = tf.reshape(self.images, [-1, input_dim, input_dim, 1])

        # First convolutional layer - maps one grayscale image to 32 feature maps.
        with tf.name_scope('conv1'):
            w_conv1 = weight_variable(
                [self.conv_size, self.conv_size, 1, self.channel_1_num])
            b_conv1 = bias_variable([self.channel_1_num])
            h_conv1 = tf.nn.relu(conv2d(x_image, w_conv1) + b_conv1)

        # Pooling layer - downsamples by 2X.
        with tf.name_scope('pool1'):
            h_pool1 = max_pool(h_conv1, self.pool_size)

        # Second convolutional layer -- maps 32 feature maps to 64.
        with tf.name_scope('conv2'):
            w_conv2 = weight_variable([self.conv_size, self.conv_size,
                                       self.channel_1_num, self.channel_2_num])
            b_conv2 = bias_variable([self.channel_2_num])
            h_conv2 = tf.nn.relu(conv2d(h_pool1, w_conv2) + b_conv2)

        # Second pooling layer.
        with tf.name_scope('pool2'):
            h_pool2 = max_pool(h_conv2, self.pool_size)

        # Fully connected layer 1 -- after 2 round of downsampling, our 28x28 image
        # is down to 7x7x64 feature maps -- maps this to 1024 features.
        last_dim = int(input_dim / (self.pool_size * self.pool_size))
        with tf.name_scope('fc1'):
            w_fc1 = weight_variable(
                [last_dim * last_dim * self.channel_2_num, self.hidden_size])
            b_fc1 = bias_variable([self.hidden_size])

        h_pool2_flat = tf.reshape(
            h_pool2, [-1, last_dim * last_dim * self.channel_2_num])
        h_fc1 = tf.nn.relu(tf.matmul(h_pool2_flat, w_fc1) + b_fc1)

        # Dropout - controls the complexity of the model, prevents co-adaptation of features.
        with tf.name_scope('dropout'):
            h_fc1_drop = tf.nn.dropout(h_fc1, self.keep_prob)

        # Map the 1024 features to 10 classes, one for each digit
        with tf.name_scope('fc2'):
            w_fc2 = weight_variable([self.hidden_size, self.y_dim])
            b_fc2 = bias_variable([self.y_dim])
            y_conv = tf.matmul(h_fc1_drop, w_fc2) + b_fc2

        with tf.name_scope('loss'):
            cross_entropy = tf.reduce_mean(
                tf.nn.softmax_cross_entropy_with_logits(labels=self.labels, logits=y_conv))
        with tf.name_scope('adam_optimizer'):
            self.train_step = tf.train.AdamOptimizer(
                self.learning_rate).minimize(cross_entropy)

        with tf.name_scope('accuracy'):
            correct_prediction = tf.equal(
                tf.argmax(y_conv, 1), tf.argmax(self.labels, 1))
            self.accuracy = tf.reduce_mean(
                tf.cast(correct_prediction, tf.float32))
498308
def build_and_train(iterations, log_stride, test=False):
  """"""

  # Test mode settings.
  batch_size = 2 if test else FLAGS.batch_size
  num_mems = 2 if test else FLAGS.num_mems
  num_heads = 1 if test else FLAGS.num_mems
  num_blocks = 1 if test else FLAGS.num_mems
  head_size = 4 if test else FLAGS.head_size
  num_objects = 2 if test else FLAGS.num_objects
  num_features = 4 if test else FLAGS.num_features
  mlp_size = (20,) if test else (256, 256, 256, 256)

  with tf.Graph().as_default():
    t0 = time.time()

    # Initialize the dataset.
    dataset = dataset_nth_farthest.NthFarthest(
        batch_size, num_objects, num_features)

    # Create the model.
    core = snt.RelationalMemory(
        mem_slots=num_mems,
        head_size=head_size,
        num_heads=num_heads,
        num_blocks=num_blocks,
        gate_style=FLAGS.gate_style)

    final_mlp = snt.nets.MLP(
        output_sizes=mlp_size,
        activate_final=True)

    model = SequenceModel(
        core=core,
        target_size=num_objects,
        final_mlp=final_mlp)

    tf.logging.info("Instantiated models ({:3f})".format(time.time() - t0))

    # Get train and test data.
    inputs_train, labels_train = dataset.get_batch()
    inputs_test, labels_test = dataset.get_batch()

    # Define target accuracy.
    def compute_accuracy(logits, targets, name="accuracy"):
      correct_pred = tf.cast(
          tf.equal(tf.cast(targets, tf.int64), tf.argmax(logits, 1)),
          tf.float32)
      return tf.reduce_mean(correct_pred, name=name)

    # Define the loss & accuracy.
    def loss_fn(inputs, labels):
      """Creates the loss and the exports."""
      logits = model(inputs)
      labels = tf.cast(labels, tf.int32)
      loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(
          logits=logits, labels=labels))
      accuracy = compute_accuracy(logits, labels)
      return loss, accuracy

    # Get training step counter.
    global_step = tf.get_variable(
        name="global_step",
        shape=[],
        dtype=tf.int64,
        initializer=tf.zeros_initializer(),
        trainable=False,
        collections=[tf.GraphKeys.GLOBAL_VARIABLES, tf.GraphKeys.GLOBAL_STEP])

    # Create the optimizer.
    learning_rate_op = tf.reduce_max([
        tf.train.exponential_decay(
            FLAGS.learning_rate,
            global_step,
            decay_steps=FLAGS.epochs // 100,
            decay_rate=0.9,
            staircase=False),
        FLAGS.min_learning_rate
    ])
    optimizer = tf.train.AdamOptimizer(learning_rate_op)
    train_loss, _ = loss_fn(inputs_train, labels_train)
    step_op = optimizer.minimize(train_loss, global_step=global_step)

    # Compute test accuracy
    logits_test = model(inputs_test)
    labels_test = tf.cast(labels_test, tf.int32)
    test_acc = compute_accuracy(logits_test, labels_test)

    tf.logging.info("Created losses and optimizers ({:3f})".format(
        time.time() - t0))

    # Begin Training.
    t0 = time.time()
    train_losses = []
    steps = []
    test_accs = []
    tf.logging.info("Starting training ({:3f})".format(time.time() - t0))
    with tf.train.SingularMonitoredSession() as sess:
      for it in six.moves.range(iterations):
        sess.run([step_op, learning_rate_op])
        if it % log_stride == 0:
          loss_v, acc_v = sess.run([train_loss, test_acc])
          elapsed = time.time() - t0
          tf.logging.info(
              "iter: {:2d}, train loss {:3f}; test acc {:3f} ({:3f})".format(
                  it, loss_v, acc_v, elapsed))
          train_losses.append(loss_v)
          steps.append(it)
        test_accs.append(acc_v)
  return steps, train_losses, test_accs
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)
787913
def create_training_instances(x):
    """"""
    (input_files, out, tokenizer, max_seq_length, dupe_factor,
     short_seq_prob, masked_lm_prob, max_predictions_per_seq, rng) = x
    time_start = time.time()
    logging.info('Processing %s', input_files)
    all_documents = [[]]

    # Input file format:
    # (1) One sentence per line. These should ideally be actual sentences, not
    # entire paragraphs or arbitrary spans of text. (Because we use the
    # sentence boundaries for the "next sentence prediction" task).
    # (2) Blank lines between documents. Document boundaries are needed so
    # that the "next sentence prediction" task doesn't span between documents.
    for input_file in input_files:
        with io.open(input_file, 'r', encoding='UTF-8') as reader:
            while True:
                line = reader.readline()
                if not line:
                    break
                line = line.strip()

                # Empty lines are used as document delimiters
                if not line:
                    all_documents.append([])
                tokens = tokenizer(line) if not args.tokenized else line.split(' ')
                if tokens:
                    all_documents[-1].append(tokens)

    # Remove empty documents
    all_documents = [x for x in all_documents if x]
    rng.shuffle(all_documents)

    vocab_words = tokenizer.vocab.idx_to_token
    instances = []
    for _ in range(dupe_factor):
        for document_index in range(len(all_documents)):
            instances.extend(
                create_instances_from_document(
                    all_documents, document_index, max_seq_length,
                    short_seq_prob, masked_lm_prob, max_predictions_per_seq,
                    vocab_words, rng))

    rng.shuffle(instances)
    input_ids = []
    segment_ids = []
    masked_lm_positions = []
    masked_lm_ids = []
    masked_lm_weights = []
    next_sentence_labels = []
    valid_lengths = []

    for inst_index, instance in enumerate(instances):
        feature = transform(instance, tokenizer, max_seq_length, max_predictions_per_seq, False)
        input_ids.append(
            np.ascontiguousarray(feature['input_ids'], dtype='int32'))
        segment_ids.append(
            np.ascontiguousarray(feature['segment_ids'], dtype='int32'))
        masked_lm_positions.append(
            np.ascontiguousarray(feature['masked_lm_positions'], dtype='int32'))
        masked_lm_ids.append(np.ascontiguousarray(feature['masked_lm_ids'], dtype='int32'))
        masked_lm_weights.append(
            np.ascontiguousarray(feature['masked_lm_weights'], dtype='float32'))
        next_sentence_labels.append(feature['next_sentence_labels'][0])
        valid_lengths.append(feature['valid_lengths'][0])
        if inst_index < 20:
            print_example(instance, feature)

    features = (input_ids, segment_ids, masked_lm_positions, masked_lm_ids, \
                masked_lm_weights, next_sentence_labels, valid_lengths)

    logging.info('*** Writing to output file %s ***', out)
    if args.format == 'numpy':
        write_to_files_np(features, tokenizer, args.max_seq_length,
                          args.max_predictions_per_seq, [out])
    elif args.format == 'recordio':
        write_to_files_rec(instances, tokenizer, args.max_seq_length,
                           args.max_predictions_per_seq, [out])
    else:
        raise ValueError('unsupported format: %s'%args.format)

    time_end = time.time()
    logging.info('Process %d files took %.1f s', len(input_files), time_end - time_start)
1104917
def build_network(self,
                      images,
                      phase_train=True,
                      nclass=1001,
                      image_depth=3,
                      data_type=tf.float32,
                      data_format="NCHW",
                      use_tf_layers=True,
                      fp16_vars=False):
        """"""
        if data_format == "NCHW":
            images = tf.transpose(images, [0, 3, 1, 2])
        var_type = tf.float32
        if data_type == tf.float16 and fp16_vars:
            var_type = tf.float16
        network = convnet_builder.ConvNetBuilder(
            images, image_depth, phase_train, use_tf_layers, data_format,
            data_type, var_type)
        with tf.variable_scope(
                "cg", custom_getter=network.get_custom_getter()):
            self.add_inference(network)
            # Add the final fully-connected class layer
            logits = (network.affine(nclass, activation="linear")
                      if not self.skip_final_affine_layer() else
                      network.top_layer)
            aux_logits = None
            if network.aux_top_layer is not None:
                with network.switch_to_aux_top_layer():
                    aux_logits = network.affine(
                        nclass, activation="linear", stddev=0.001)
        if data_type == tf.float16:
            # TODO(reedwm): Determine if we should do this cast here.
            logits = tf.cast(logits, tf.float32)
            if aux_logits is not None:
                aux_logits = tf.cast(aux_logits, tf.float32)
        return logits, aux_logits
1010523
def build_and_train(iterations, log_stride, test=False):
  """"""

  # Test mode settings.
  batch_size = 2 if test else FLAGS.batch_size
  num_mems = 2 if test else FLAGS.num_mems
  num_heads = 1 if test else FLAGS.num_mems
  num_blocks = 1 if test else FLAGS.num_mems
  head_size = 4 if test else FLAGS.head_size
  max_length = 3 if test else FLAGS.max_length
  max_nest = 2 if test else FLAGS.max_nest
  mlp_size = (20,) if test else (256, 256, 256, 256)

  with tf.Graph().as_default():
    t0 = time.time()

    # Initialize the dataset.
    lte_train = learn_to_execute.LearnToExecute(
        batch_size, max_length, max_nest)
    lte_test = learn_to_execute.LearnToExecute(
        batch_size, max_length, max_nest, mode=learn_to_execute.Mode.TEST)
    train_data_iter = lte_train.make_one_shot_iterator().get_next()
    test_data_iter = lte_test.make_one_shot_iterator().get_next()
    output_size = lte_train.state.vocab_size

    # Create the model.
    core = snt.RelationalMemory(
        mem_slots=num_mems,
        head_size=head_size,
        num_heads=num_heads,
        num_blocks=num_blocks,
        gate_style=FLAGS.gate_style)
    final_mlp = snt.nets.MLP(
        output_sizes=mlp_size,
        activate_final=True)
    model = SequenceModel(
        core=core,
        target_size=output_size,
        final_mlp=final_mlp)
    tf.logging.info("Instantiated models ({:3f})".format(time.time() - t0))

    # Define the loss & accuracy.
    def loss_fn(inputs, targets, input_sequence_length, output_sequence_length):
      """Creates the loss and the exports."""
      logits = model(
          inputs, targets, input_sequence_length, output_sequence_length)
      targets = tf.cast(targets, tf.int32)
      sq_sz_out_max = targets.shape[0].value

      # Create a mask to ignore accuracy on buffer characters.
      sequence_sizes = tf.cast(output_sequence_length, tf.float32)
      lengths_transposed = tf.expand_dims(sequence_sizes, 1)
      range_row = tf.expand_dims(
          tf.range(0, sq_sz_out_max, 1, dtype=tf.float32), 0)
      mask = tf.cast(tf.transpose(tf.less(range_row, lengths_transposed)),
                     tf.float32)

      # Compute token accuracy and solved.
      correct = tf.equal(tf.argmax(logits, 2), tf.argmax(targets, 2))
      solved = tf.reduce_all(tf.boolean_mask(correct, tf.squeeze(mask)), axis=0)
      token_acc = tf.reduce_sum(tf.cast(correct, tf.float32) * mask)
      token_acc /= tf.reduce_sum(sequence_sizes)

      # Compute Loss.
      mask = tf.cast(tf.tile(tf.expand_dims(mask, 2), (1, 1, logits.shape[2])),
                     tf.float32)
      masked_logits = logits * mask
      masked_target = tf.cast(targets, tf.float32) * mask
      logits_flat = tf.reshape(masked_logits,
                               [sq_sz_out_max * batch_size, -1])
      target_flat = tf.reshape(masked_target,
                               [sq_sz_out_max * batch_size, -1])
      xent = tf.nn.softmax_cross_entropy_with_logits(logits=logits_flat,
                                                     labels=target_flat)
      loss = tf.reduce_mean(xent)
      return loss, token_acc, solved

    # Get training step counter.
    global_step = tf.train.get_or_create_global_step()

    # Create the optimizer.
    learning_rate_op = tf.reduce_max([
        tf.train.exponential_decay(
            FLAGS.learning_rate,
            global_step,
            decay_steps=FLAGS.epochs // 100,
            decay_rate=0.9,
            staircase=False),
        FLAGS.min_learning_rate
    ])
    optimizer = tf.train.AdamOptimizer(learning_rate_op)

    # Compute loss, accuracy & the step op.
    inputs, targets, _, input_lengths, output_lengths = train_data_iter
    train_loss, train_acc, train_sol = loss_fn(
        inputs, targets, input_lengths, output_lengths)
    step_op = optimizer.minimize(train_loss, global_step=global_step)
    inputs, targets, _, input_lengths, output_lengths = test_data_iter
    _, test_acc, test_sol = loss_fn(
        inputs, targets, input_lengths, output_lengths)
    tf.logging.info("Created losses and optimizers ({:3f})".format(
        time.time() - t0))

    # Begin Training.
    t0 = time.time()
    tf.logging.info("Starting training ({:3f})".format(time.time() - t0))
    with tf.train.SingularMonitoredSession() as sess:
      for it in six.moves.range(iterations):
        sess.run([step_op, learning_rate_op])
        if it % log_stride == 0:
          loss_v, train_acc_v, test_acc_v, train_sol_v, test_sol_v = sess.run([
              train_loss, train_acc, test_acc, train_sol, test_sol])
          elapsed = time.time() - t0
          tf.logging.info(
              "iter: {:2d}, train loss {:3f}; train acc {:3f}; test acc {:3f};"
              " train solved {:3f}; test solved {:3f}; ({:3f})".format(
                  it, loss_v, train_acc_v, test_acc_v, train_sol_v, test_sol_v,
                  elapsed))
144711
def build_graph(self, image, label):
        """"""

        # In tensorflow, inputs to convolution function are assumed to be
        # NHWC. Add a single channel here.
        image = tf.expand_dims(image, 3)

        image = image * 2 - 1   # center the pixels values at zero

        # The context manager `argscope` sets the default option for all the layers under
        # this context. Here we use 32 channel convolution with shape 3x3
        with argscope([tf.layers.conv2d], padding='same', activation=tf.nn.relu):
            l = tf.layers.conv2d(image, 32, 3, name='conv0')
            l = tf.layers.max_pooling2d(l, 2, 2, padding='valid')
            l = tf.layers.conv2d(l, 32, 3, name='conv1')
            l = tf.layers.conv2d(l, 32, 3, name='conv2')
            l = tf.layers.max_pooling2d(l, 2, 2, padding='valid')
            l = tf.layers.conv2d(l, 32, 3, name='conv3')
            l = tf.layers.flatten(l)
            l = tf.layers.dense(l, 512, activation=tf.nn.relu, name='fc0')
            l = tf.layers.dropout(l, rate=0.5,
                                  training=get_current_tower_context().is_training)
        logits = tf.layers.dense(l, 10, activation=tf.identity, name='fc1')

        # a vector of length B with loss of each sample
        cost = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=label)
        cost = tf.reduce_mean(cost, name='cross_entropy_loss')  # the average cross-entropy loss

        correct = tf.cast(tf.nn.in_top_k(logits, label, 1), tf.float32, name='correct')
        accuracy = tf.reduce_mean(correct, name='accuracy')

        # This will monitor training error & accuracy (in a moving average fashion). The value will be automatically
        # 1. written to tensosrboard
        # 2. written to stat.json
        # 3. printed after each epoch
        train_error = tf.reduce_mean(1 - correct, name='train_error')
        summary.add_moving_summary(train_error, accuracy)

        # Use a regex to find parameters to apply weight decay.
        # Here we apply a weight decay on all W (weight matrix) of all fc layers
        # If you don't like regex, you can certainly define the cost in any other methods.
        wd_cost = tf.multiply(1e-5,
                              regularize_cost('fc.*/kernel', tf.nn.l2_loss),
                              name='regularize_loss')
        total_cost = tf.add_n([wd_cost, cost], name='total_cost')
        summary.add_moving_summary(cost, wd_cost, total_cost)

        # monitor histogram of all weight (of conv and fc layers) in tensorboard
        summary.add_param_summary(('.*/kernel', ['histogram', 'rms']))
        # the function should return the total cost to be optimized
        return total_cost
692314
def from_learner(cls, learn: Learner,  ds_type:DatasetType=DatasetType.Valid):
        ""
        preds = learn.get_preds(ds_type=ds_type, with_loss=True)
        return cls(learn, *preds)
722209
def general_conv(x,
                 num_filters=64,
                 filter_size=7,
                 stride=1,
                 stddev=0.02,
                 padding="VALID",
                 name="conv",
                 do_norm="instance",
                 do_relu=True,
                 relufactor=0):
  """"""
  with tf.variable_scope(name):
    x = layers().Conv2D(
        num_filters,
        filter_size,
        stride,
        padding,
        activation=None,
        kernel_initializer=tf.truncated_normal_initializer(stddev=stddev),
        bias_initializer=tf.constant_initializer(0.0))(x)
    if do_norm == "layer":
      x = layer_norm(x)
    elif do_norm == "instance":
      x = instance_norm(x)

    if do_relu:
      if relufactor == 0:
        x = tf.nn.relu(x, "relu")
      else:
        x = lrelu(x, leak=relufactor)

    return x
1341226
def train(data_train, model, nsp_loss, mlm_loss, vocab_size, ctx, store):
    """"""
    mlm_metric = nlp.metric.MaskedAccuracy()
    nsp_metric = nlp.metric.MaskedAccuracy()
    mlm_metric.reset()
    nsp_metric.reset()

    lr = args.lr
    optim_params = {'learning_rate': lr, 'epsilon': 1e-6, 'wd': 0.01}
    if args.dtype == 'float16':
        optim_params['multi_precision'] = True

    trainer = mx.gluon.Trainer(model.collect_params(), 'bertadam', optim_params,
                               update_on_kvstore=False, kvstore=store)
    dynamic_loss_scale = args.dtype == 'float16'
    fp16_trainer = FP16Trainer(trainer, dynamic_loss_scale=dynamic_loss_scale)

    if args.ckpt_dir and args.start_step:
        state_path = os.path.join(args.ckpt_dir, '%07d.states' % args.start_step)
        logging.info('Loading trainer state from %s', state_path)
        trainer.load_states(state_path)

    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 = running_num_tks = 0
    batch_num = 0
    step_num = args.start_step

    parallel_model = ParallelBERT(model, mlm_loss, nsp_loss, vocab_size,
                                  store.num_workers * accumulate, trainer=fp16_trainer)
    num_ctxes = len(ctx)
    parallel = nlp.utils.Parallel(num_ctxes if num_ctxes > 1 else 0, parallel_model)

    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*num_ctxes, 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, 12, profile_name=args.profile)
                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:
                    if data_batch[0].shape[0] < len(ctx):
                        continue
                    data_list = split_and_load(data_batch, ctx)

                ns_label_list, ns_pred_list = [], []
                mask_label_list, mask_pred_list, mask_weight_list = [], [], []

                # parallel forward / backward
                for data in data_list:
                    parallel.put(data)
                for _ in range(len(ctx)):
                    (_, next_sentence_label, classified, masked_id,
                     decoded, masked_weight, ls1, ls2, valid_length) = parallel.get()
                    ns_label_list.append(next_sentence_label)
                    ns_pred_list.append(classified)
                    mask_label_list.append(masked_id)
                    mask_pred_list.append(decoded)
                    mask_weight_list.append(masked_weight)
                    running_mlm_loss += ls1.as_in_context(mx.cpu()) / num_ctxes
                    running_nsp_loss += ls2.as_in_context(mx.cpu()) / num_ctxes
                    running_num_tks += valid_length.sum().as_in_context(mx.cpu())

                # update
                if (batch_num + 1) % accumulate == 0:
                    fp16_trainer.step(1, max_norm=1)
                nsp_metric.update(ns_label_list, ns_pred_list)
                mlm_metric.update(mask_label_list, mask_pred_list, mask_weight_list)
                # 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:
                    save_params(step_num, model, trainer, args.ckpt_dir)
                batch_num += 1
    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))
1212792
def unet_learner(data:DataBunch, arch:Callable, pretrained:bool=True, blur_final:bool=True,
                 norm_type:Optional[NormType]=NormType, split_on:Optional[SplitFuncOrIdxList]=None, blur:bool=False,
                 self_attention:bool=False, y_range:Optional[Tuple[float,float]]=None, last_cross:bool=True,
                 bottle:bool=False, cut:Union[int,Callable]=None, **learn_kwargs:Any)->Learner:
    ""
    meta = cnn_config(arch)
    body = create_body(arch, pretrained, cut)
    model = to_device(models.unet.DynamicUnet(body, n_classes=data.c, blur=blur, blur_final=blur_final,
          self_attention=self_attention, y_range=y_range, norm_type=norm_type, last_cross=last_cross,
          bottle=bottle), data.device)
    learn = Learner(data, model, **learn_kwargs)
    learn.split(ifnone(split_on, meta['split']))
    if pretrained: learn.freeze()
    apply_init(model[2], nn.init.kaiming_normal_)
    return learn
606573
def deep_discriminator(x,
                       batch_norm,
                       is_training,
                       filters=64,
                       filter_size=4,
                       stride=2,
                       output_size=1024):
  """"""
  with tf.variable_scope(
      "discriminator", initializer=tf.random_normal_initializer(stddev=0.02)):
    batch_size, height, width = shape_list(x)[:3]  # pylint: disable=unbalanced-tuple-unpacking
    net = layers().Conv2D(
        filters, filter_size, strides=stride, padding="SAME", name="conv1")(x)
    net = lrelu(net)
    net = layers().Conv2D(
        2 * filters,
        filter_size,
        strides=stride,
        padding="SAME",
        name="conv2")(net)
    # [bs, h/4, w/4, 128]
    if batch_norm:
      net = layers().BatchNormalization(
          training=is_training, momentum=0.999, name="d_bn2")(net)
    net = lrelu(net)
    size = height * width
    x_shape = x.get_shape().as_list()
    if x_shape[1] is None or x_shape[2] is None:
      net = tf.reduce_mean(net, axis=[1, 2])
    else:
      net = tf.reshape(net, [batch_size, size * 8])
    net = layers().Dense(output_size, name="d_fc3")(net)
    if batch_norm:
      net = layers().BatchNormalization(
          training=is_training, momentum=0.999, name="d_bn3")(net)
    net = lrelu(net)
    return net
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
85522
def generator(self, output, target):
        ""
        fake_pred = self.gan_model.critic(output)
        return self.loss_funcG(fake_pred, target, output)
1362988
def learn(self, text):
        """"""
        if type(text) != types.UnicodeType:
            # Assume that non-Unicode text is encoded as utf-8, which
            # should be somewhat safe in the modern world.
            text = text.decode("utf-8", "ignore")

        tokens = self.tokenizer.split(text)
        trace("Brain.learn_input_token_count", len(tokens))

        self._learn_tokens(tokens)
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))