Construct the data, model, loss and optimizer then train.
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
738660
def _init_model(self):
        
        base_model = self._get_base_model()

        top_layer_model = base_model.output
        top_layer_model = GlobalAveragePooling2D()(top_layer_model)
        top_layer_model = Dense(1024, activation='relu')(top_layer_model)
        prediction_layer = Dense(output_dim=len(self.emotion_map.keys()), activation='softmax')(top_layer_model)

        model = Model(input=base_model.input, output=prediction_layer)
        print(model.summary())
        for layer in base_model.layers:
            layer.trainable = False
        model.compile(optimizer='rmsprop', loss='categorical_crossentropy', metrics=['accuracy'])

        self.model = model
667424
def train(self, params):
    
    if params['reset_every_training']:
      n = params['encoding_num']
      self.net = buildNetwork(n, params['num_cells'], n,
                               hiddenclass=LSTMLayer,
                               bias=True,
                               outputbias=params['output_bias'],
                               recurrent=True)
      self.net.reset()

    # prepare training dataset
    ds = SequentialDataSet(params['encoding_num'], params['encoding_num'])
    history = self.window(self.history, params)
    resets = self.window(self.resets, params)

    for i in xrange(1, len(history)):
      if not resets[i - 1]:
        ds.addSample(self.encoder.encode(history[i - 1]),
                     self.encoder.encode(history[i]))
      if resets[i]:
        ds.newSequence()

    print "Train LSTM network on buffered dataset of length ", len(history)
    if params['num_epochs'] > 1:
      trainer = RPropMinusTrainer(self.net,
                                  dataset=ds,
                                  verbose=params['verbosity'] > 0)

      if len(history) > 1:
        trainer.trainEpochs(params['num_epochs'])

      # run network on buffered dataset after training to get the state right
      self.net.reset()
      for i in xrange(len(history) - 1):
        symbol = history[i]
        output = self.net.activate(self.encoder.encode(symbol))
        self.encoder.classify(output, num=params['num_predictions'])

        if resets[i]:
          self.net.reset()
    else:
      self.trainer.setData(ds)
      self.trainer.train()

      # run network on buffered dataset after training to get the state right
      self.net.reset()
      for i in xrange(len(history) - 1):
        symbol = history[i]
        output = self.net.activate(self.encoder.encode(symbol))
        self.encoder.classify(output, num=params['num_predictions'])

        if resets[i]:
          self.net.reset()
232983
def train_model(self, balance, *args, **kwargs):
        
        self.balance = np.array(balance)
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'))
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)
1006269
def train(self, data_iterator):
        
        optimizer = get_optimizer(self.master_optimizer)
        self.model = model_from_yaml(self.yaml, self.custom_objects)
        self.model.compile(optimizer=optimizer,
                           loss=self.master_loss, metrics=self.master_metrics)
        self.model.set_weights(self.parameters.value)

        feature_iterator, label_iterator = tee(data_iterator, 2)
        x_train = np.asarray([x for x, y in feature_iterator])
        y_train = np.asarray([y for x, y in label_iterator])

        self.model.compile(optimizer=self.master_optimizer,
                           loss=self.master_loss,
                           metrics=self.master_metrics)

        weights_before_training = self.model.get_weights()
        if x_train.shape[0] > self.train_config.get('batch_size'):
            self.model.fit(x_train, y_train, **self.train_config)
        weights_after_training = self.model.get_weights()
        deltas = subtract_params(
            weights_before_training, weights_after_training)
        yield deltas
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)
500802
def main(args = None):
    '''
    
    '''

    parser = argparse.ArgumentParser(description = 'Tool to perform transfer learning')

    parser.add_argument('-c','--configure',
                        action = 'store_true',
                        help = 'Configure transfer')

    parser.add_argument('-e','--export',
                        action = 'store_true',
                        dest = 'export_config',
                        help = 'Export configuration and models')

    parser.add_argument('-i','--import',
                        action = 'store',
                        type = str,
                        default = None,
                        dest = 'import_config',
                        help = 'Import configuration and models')

    parser.add_argument('-p','--project',
                        action = 'store',
                        type = str,
                        default = None,
                        dest = 'project',
                        help = 'Specify a project, if not supplied it will be picked from configured projects')

    parser.add_argument('-r','--run',
                        action = 'store_true',
                        help = 'Run all transfer learning operations')

    parser.add_argument('-f','--final',
                        action = 'store_true',
                        help = 'Run final training on all layers: Warning SLOW!')

    parser.add_argument('-l','--last-weights',
                        action = 'store_true',
                        dest = 'last',
                        help = 'Restart from the last weights, rather than the best intermediate weights')

    parser.add_argument('--predict',
                        action = 'store',
                        type = str,
                        default = None,
                        const = 'default',
                        dest = 'predict',
                        nargs='?',
                        help = 'Predict model on file or directory')

    parser.add_argument('--prediction-rest-api',
                        action = 'store_true',
                        dest = 'rest_api',
                        help = 'Start rest api to make predictions on files or directories')

    if len(sys.argv) == 1:
        parser.print_help()
        return

    args = parser.parse_args()

    if args.import_config is not None:
        import_config(args.import_config)
        return
    elif args.export_config:
        project = select_project(args.project)
        weights = model_input(project)
        ind = model_individual_input(project, weights)
        export_config(project, weights, ind)
        return
    elif args.configure:
        configure()
        return
    else:
        project = select_project(args.project)

    if args.run:

        if project['is_array'] == False:
            project = images_to_array(project)
            update_config(project)

        if project['is_augmented'] == False:
            project = augment_arrays(project)
            update_config(project)

        if project['is_pre_model'] == False:
            project = pre_model(project)
            update_config(project)

        project = train_model(project, final = args.final, last = args.last)
        update_config(project)

        print('')
        print(colored('Completed modeling round: ' + str(project['model_round']), 'cyan'))
        print('')
        print('Best current model: ', colored(project['best_weights'], 'yellow'))
        print('Last current model: ', colored(project['last_weights'], 'yellow'))
        print('')
        print('To further refine the model, run again with:')
        print('')
        print(colored('    transfer --run --project ' + project['name'], 'green'))
        print('')
        print('To make predictions:')
        print('')
        print(colored('    transfer --predict [optional dir or file] --project ' + project['name'], 'yellow'))
        print('')

    elif args.rest_api:
        if project['server_weights'] is not None:
            start_server(project, 'server_weights')

        elif project['best_weights'] is not None:
            weights = model_input(project)
            start_server(project, weights)

        else:
            print('Model is not trained.  Please first run your project:')
            print('')
            print(colored('    transfer --run', 'green'))
            print('')
    elif args.predict is not None:
        if args.predict == 'default':

            completer = Completer()
            readline.set_completer_delims('\t')
            readline.parse_and_bind('tab: complete')
            readline.set_completer(completer.path_completer)
            args.predict = str_input('Enter a path to file(s): ')
        if project['server_weights'] is not None:
            predict_model(project, 'server_weights', args.predict)

        elif project['best_weights'] is not None:
            weights = model_input(project)
            print('Predicting on image(s) in: ', colored(args.predict, 'yellow'))
            predict_model(project, weights, args.predict)

        else:
            print('Model is not trained.  Please first run your project:')
            print('')
            print(colored('    transfer --run', 'green'))
            print('')
1095263
def build_model(self, n_features, n_classes):
        
        self._create_placeholders(n_features, n_classes)
        self._create_variables(n_features)

        next_train = self._create_encoding_layers()
        self.mod_y, _, _ = Layers.linear(next_train, n_classes)
        self.layer_nodes.append(self.mod_y)

        self.cost = self.loss.compile(self.mod_y, self.input_labels)
        self.train_step = self.trainer.compile(self.cost)
        self.accuracy = Evaluation.accuracy(self.mod_y, self.input_labels)
394945
def compile(self, optimizer, loss, metrics=None):
        
        if isinstance(loss, six.string_types):
            loss = [loss]
        if metrics is None:
            metrics = []
        if isinstance(metrics, six.string_types):
            metrics = [metrics]

        self._stats_to_inference = loss + metrics + [TOTAL_LOSS_NAME]
        setup_keras_trainer(
            self.trainer, get_model=self.get_model,
            input_signature=self.input_signature,
            target_signature=self.target_signature,
            input=self.input,
            optimizer=optimizer,
            loss=loss,
            metrics=metrics)
692834
def train(self):
        ''''''

        # Initialize the train subcommand's argparser
        parser = argparse.ArgumentParser(description='Train a dialogue model on a dialogue corpus or a dsrt dataset')
        self.init_train_args(parser)

        # Parse the args we got
        args = parser.parse_args(sys.argv[2:])
        args.config = ConfigurationLoader(args.config).load().model_config

        print(CLI_DIVIDER + '\n')
        Trainer(**vars(args)).run()
1148745
def main(model=None, output_dir=None, n_iter=100):
    """"""
    if model is not None:
        nlp = spacy.load(model)  # load existing spaCy model
        print("Loaded model '%s'" % model)
    else:
        nlp = spacy.blank("en")  # create blank Language class
        print("Created blank 'en' model")

    # create the built-in pipeline components and add them to the pipeline
    # nlp.create_pipe works for built-ins that are registered with spaCy
    if "ner" not in nlp.pipe_names:
        ner = nlp.create_pipe("ner")
        nlp.add_pipe(ner, last=True)
    # otherwise, get it so we can add labels
    else:
        ner = nlp.get_pipe("ner")

    # add labels
    for _, annotations in TRAIN_DATA:
        for ent in annotations.get("entities"):
            ner.add_label(ent[2])

    # get names of other pipes to disable them during training
    other_pipes = [pipe for pipe in nlp.pipe_names if pipe != "ner"]
    with nlp.disable_pipes(*other_pipes):  # only train NER
        # reset and initialize the weights randomly – but only if we're
        # training a new model
        if model is None:
            nlp.begin_training()
        for itn in range(n_iter):
            random.shuffle(TRAIN_DATA)
            losses = {}
            # batch up the examples using spaCy's minibatch
            batches = minibatch(TRAIN_DATA, size=compounding(4.0, 32.0, 1.001))
            for batch in batches:
                texts, annotations = zip(*batch)
                nlp.update(
                    texts,  # batch of texts
                    annotations,  # batch of annotations
                    drop=0.5,  # dropout - make it harder to memorise data
                    losses=losses,
                )
            print("Losses", losses)

    # test the trained model
    for text, _ in TRAIN_DATA:
        doc = nlp(text)
        print("Entities", [(ent.text, ent.label_) for ent in doc.ents])
        print("Tokens", [(t.text, t.ent_type_, t.ent_iob) for t in doc])

    # save model to output directory
    if output_dir is not None:
        output_dir = Path(output_dir)
        if not output_dir.exists():
            output_dir.mkdir()
        nlp.to_disk(output_dir)
        print("Saved model to", output_dir)

        # test the saved model
        print("Loading from", output_dir)
        nlp2 = spacy.load(output_dir)
        for text, _ in TRAIN_DATA:
            doc = nlp2(text)
            print("Entities", [(ent.text, ent.label_) for ent in doc.ents])
            print("Tokens", [(t.text, t.ent_type_, t.ent_iob) for t in doc])
1048124
def main(model_folder):
    """"""
    model_description_file = os.path.join(model_folder, "info.yml")

    # Read the model description file
    with open(model_description_file, 'r') as ymlfile:
        model_description = yaml.load(ymlfile)

    # Analyze model
    logging.info(model_description['model'])
    data = {}
    data['training'] = os.path.join(model_folder, "traindata.hdf5")
    data['testing'] = os.path.join(model_folder, "testdata.hdf5")
    data['validating'] = os.path.join(model_folder, "validdata.hdf5")
    train_model(model_folder)
899011
def build_iters(data_dir, max_records, q, horizon, splits, batch_size):
    
    # Read in data as numpy array
    df = pd.read_csv(os.path.join(data_dir, "electricity.txt"), sep=",", header=None)
    feature_df = df.iloc[:, :].astype(float)
    x = feature_df.as_matrix()
    x = x[:max_records] if max_records else x

    # Construct training examples based on horizon and window
    x_ts = np.zeros((x.shape[0] - q, q, x.shape[1]))
    y_ts = np.zeros((x.shape[0] - q, x.shape[1]))
    for n in range(x.shape[0]):
        if n + 1 < q:
            continue
        elif n + 1 + horizon > x.shape[0]:
            continue
        else:
            y_n = x[n + horizon, :]
            x_n = x[n + 1 - q:n + 1, :]
        x_ts[n-q] = x_n
        y_ts[n-q] = y_n

    # Split into training and testing data
    training_examples = int(x_ts.shape[0] * splits[0])
    valid_examples = int(x_ts.shape[0] * splits[1])
    x_train, y_train = x_ts[:training_examples], \
                       y_ts[:training_examples]
    x_valid, y_valid = x_ts[training_examples:training_examples + valid_examples], \
                       y_ts[training_examples:training_examples + valid_examples]
    x_test, y_test = x_ts[training_examples + valid_examples:], \
                     y_ts[training_examples + valid_examples:]

    #build iterators to feed batches to network
    train_iter = mx.io.NDArrayIter(data=x_train,
                                   label=y_train,
                                   batch_size=batch_size)
    val_iter = mx.io.NDArrayIter(data=x_valid,
                                 label=y_valid,
                                 batch_size=batch_size)
    test_iter = mx.io.NDArrayIter(data=x_test,
                                  label=y_test,
                                  batch_size=batch_size)
    return train_iter, val_iter, test_iter
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.")
44535
def train_and_eval(train_batch_size, test_batch_size, num_hidden, learning_rate,
                   num_train_steps, report_every, test_every):
  """"""

  data_dict = dataset_mnist.get_data("mnist", train_batch_size, test_batch_size)
  train_data = data_dict["train_iterator"]
  test_data = data_dict["test_iterator"]

  # Sonnet separates the configuration of a model from its attachment into the
  # graph. Here we configure the shape of the model, but this call does not
  # place any ops into the graph.
  mlp = snt.nets.MLP([num_hidden, data_dict["num_classes"]])

  train_images, train_labels = train_data.get_next()
  test_images, test_labels = test_data.get_next()

  # Flatten images to pass to model.
  train_images = snt.BatchFlatten()(train_images)
  test_images = snt.BatchFlatten()(test_images)

  # Call our model, which creates it in the graph. Our build function
  # is parameterized by the source of images, and here we connect the model to
  # the training images.
  train_logits = mlp(train_images)

  # Training loss and optimizer.
  loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
      labels=train_labels, logits=train_logits)
  loss_avg = tf.reduce_mean(loss)
  optimizer = tf.train.GradientDescentOptimizer(learning_rate)
  optimizer_step = optimizer.minimize(loss_avg)

  # As before, we make a second instance of our model in the graph, which shares
  # its parameters with the first instance of the model. The Sonnet Module code
  # takes care of the variable sharing for us: because we are calling the same
  # instance of Model, we will automatically reference the same, shared
  # variables.
  test_logits = mlp(test_images)
  test_classes = tf.nn.softmax(test_logits)
  test_correct = tf.nn.in_top_k(test_classes, test_labels, k=1)

  with tf.train.SingularMonitoredSession() as sess:

    for step_idx in range(num_train_steps):
      current_loss, _ = sess.run([loss_avg, optimizer_step])
      if step_idx % report_every == 0:
        tf.logging.info("Step: %4d of %d - loss: %.02f.",
                        step_idx + 1, num_train_steps, current_loss)
      if step_idx % test_every == 0:
        sess.run(test_data.initializer)
        current_correct = sess.run(test_correct)
        correct_count = np.count_nonzero(current_correct)
        tf.logging.info("Test: %d of %d correct.",
                        correct_count, test_batch_size)