jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def build_model(self):

        self.X = tf.placeholder(tf.float32, [None, self.dims[0]])

        self.Y = tf.placeholder(tf.float32, [None, self.max_length])

        self.N_sample = tf.placeholder(dtype=tf.float32, shape=[1])


        lstm_cell = tf.contrib.rnn.BasicLSTMCell(num_units=self.rnn_hidden_size)

        self.segment_ids = tf.placeholder(tf.float32, [self.max_length])

        self.rnn_dim_feature = np.shape(self.data_X)[-1]

        self.rnn_input = tf.constant(value=self.data_X, dtype=tf.float32, shape=[self.num_example,self.max_length, self.rnn_dim_feature]) #

        state = lstm_cell.zero_state(self.num_example, dtype=tf.float32)

        xx = tf.unstack(self.rnn_input, self.max_length, 1)

        rnn_output, state = rnn.static_rnn(lstm_cell, xx, initial_state=state)

        self.train_loss = 0
        self.YY = {}

        for step in range(self.max_length):


            self.YY[str(step)] = tf.slice(self.Y, [0, step], [-1, 1])

            ss = tf.slice(rnn_output, [step, 0, 0], [1, -1, -1])

            self.embed[str(step)] = self.deepset(tf.squeeze(ss))


            state_embedding = tf.tile(self.embed[str(step)], [tf.cast(self.N_sample[0], tf.int32), 1])



            self.bo_net(self.X,state_embedding, step)


            self.train_loss += self.ml_primal[str(step)]




        self.opt = tf.train.AdamOptimizer(learning_rate=self.lr).minimize(self.train_loss)
//  [python]
def _get_global_attn_indices(is_index_global_attn):
        """ compute global attn indices required throughout forward pass """
        # helper variable
        num_global_attn_indices = tf.math.count_nonzero(is_index_global_attn, axis=1)
        num_global_attn_indices = tf.cast(num_global_attn_indices, dtype=tf.constant(1).dtype)

        # max number of global attn indices in batch
        max_num_global_attn_indices = tf.reduce_max(num_global_attn_indices)

        # indices of global attn
        is_index_global_attn_nonzero = tf.where(is_index_global_attn)

        # helper variable
        is_local_index_global_attn = tf.range(max_num_global_attn_indices) < tf.expand_dims(
            num_global_attn_indices, axis=-1
        )

        # location of the non-padding values within global attention indices
        is_local_index_global_attn_nonzero = tf.where(is_local_index_global_attn)

        # location of the padding values within global attention indices
        is_local_index_no_global_attn_nonzero = tf.where(tf.math.logical_not(is_local_index_global_attn))

        return (
            max_num_global_attn_indices,
            is_index_global_attn_nonzero,
            is_local_index_global_attn_nonzero,
            is_local_index_no_global_attn_nonzero,
        )
//  [python]
def decode_box_outputs(pred_boxes, anchor_boxes):
  """Transforms relative regression coordinates to absolute positions.

  Network predictions are normalized and relative to a given anchor; this
  reverses the transformation and outputs absolute coordinates for the input
  image.

  Args:
    pred_boxes: predicted box regression targets.
    anchor_boxes: anchors on all feature levels.
  Returns:
    outputs: bounding boxes.
  """
  anchor_boxes = t.cast(anchor_boxes, pred_boxes.dtype)
  ycenter_a = (anchor_boxes[..., 0] + anchor_boxes[..., 2]) / 2
  xcenter_a = (anchor_boxes[..., 1] + anchor_boxes[..., 3]) / 2
  ha = anchor_boxes[..., 2] - anchor_boxes[..., 0]
  wa = anchor_boxes[..., 3] - anchor_boxes[..., 1]
  ty, tx, th, tw = tf.unstack(pred_boxes, num=4, axis=-1)

  w = tf.math.exp(tw) * wa
  h = tf.math.exp(th) * ha
  ycenter = ty * ha + ycenter_a
  xcenter = tx * wa + xcenter_a
  ymin = ycenter - h / 2.
  xmin = xcenter - w / 2.
  ymax = ycenter + h / 2.
  xmax = xcenter + w / 2.
  return tf.stack([ymin, xmin, ymax, xmax], axis=-1)
//  [python]
def decode_box_outputs(pred_boxes, anchor_boxes):
  """Transforms relative regression coordinates to absolute positions.

  Network predictions are normalized and relative to a given anchor; this
  reverses the transformation and outputs absolute coordinates for the input
  image.

  Args:
    pred_boxes: predicted box regression targets.
    anchor_boxes: anchors on all feature levels.
  Returns:
    outputs: bounding boxes.
  """
  anchor_boxes = tf.cast(anchor_boxes, pred_boxes.dtype)
  ycenter_a = (anchor_boxes[..., 0] + anchor_boxes[..., 2]) / 2
  xcenter_a = (anchor_boxes[..., 1] + anchor_boxes[..., 3]) / 2
  ha = anchor_boxes[..., 2] - anchor_boxes[..., 0]
  wa = anchor_boxes[..., 3] - anchor_boxes[..., 1]
  ty, tx, th, tw = tf.unstack(pred_boxes, num=4, axis=-1)

  w = tf.math.exp(tw) * wa
  h = tf.math.exp(th) * ha
  ycenter = ty * ha + ycenter_a
  xcenter = tx * wa + xcenter_a
  ymin = ycenter - h / 2.
  xmin = xcenter - w / 2.
  ymax = ycenter + h / 2.
  xmax = xcenter + w / 2.
  return tf.stack([ymin, xmin, ymax, xmax], axis=-1)
//  [python]
def _fake_build_inputs(self, params, input_context=None):  # pylint: disable=unused-argument
  def dummy_data(_):
    dummy_ids = tf.zeros((1, params.seq_length), dtype=tf.int32)
    dummy_lm = tf.zeros((1, params.max_predictions_per_seq), dtype=tf.int32)
    return dict(
        input_word_ids=dummy_ids,
        input_mask=dummy_ids,
        masked_lm_positions=dummy_lm,
        masked_lm_ids=dummy_lm,
        masked_lm_weights=tf.cast(dummy_lm, dtype=tf.float32))

  dataset = tf.data.Dataset.range(1)
  dataset = dataset.repeat()
  dataset = dataset.map(
      dummy_data, num_parallel_calls=tf.data.experimental.AUTOTUNE)
  return dataset
//  [python]
def label_anchors(self, gt_boxes, gt_labels):
    """Labels anchors with ground truth inputs.

    Args:
      gt_boxes: A float tensor with shape [N, 4] representing groundtruth boxes.
        For each row, it stores [y0, x0, y1, x1] for four corners of a box.
      gt_labels: A integer tensor with shape [N, 1] representing groundtruth
        classes.
    Returns:
      cls_targets_dict: ordered dictionary with keys
        [min_level, min_level+1, ..., max_level]. The values are tensor with
        shape [height_l, width_l, num_anchors]. The height_l and width_l
        represent the dimension of class logits at l-th level.
      box_targets_dict: ordered dictionary with keys
        [min_level, min_level+1, ..., max_level]. The values are tensor with
        shape [height_l, width_l, num_anchors * 4]. The height_l and
        width_l represent the dimension of bounding box regression output at
        l-th level.
      num_positives: scalar tensor storing number of positives in an image.
    """
    gt_box_list = box_list.BoxList(gt_boxes)
    anchor_box_list = box_list.BoxList(self._anchors.boxes)

    # cls_weights, box_weights are not used
    cls_targets, _, box_targets, _, matches = self._target_assigner.assign(
        anchor_box_list, gt_box_list, gt_labels)

    # class labels start from 1 and the background class = -1
    cls_targets -= 1
    cls_targets = t.cast(cls_targets, t.int32)

    # Unpack labels.
    cls_targets_dict = self._unpack_labels(cls_targets)
    box_targets_dict = self._unpack_labels(box_targets)
    num_positives = t.reduce_sum(
        t.cast(t.not_equal(matches.match_results, -1), t.float32))

    return cls_targets_dict, box_targets_dict, num_positives
//  [python]
def label_anchors(self, gt_boxes, gt_labels):
    """Labels anchors with ground truth inputs.

    Args:
      gt_boxes: A float tensor with shape [N, 4] representing groundtruth boxes.
        For each row, it stores [y0, x0, y1, x1] for four corners of a box.
      gt_labels: A integer tensor with shape [N, 1] representing groundtruth
        classes.
    Returns:
      cls_targets_dict: ordered dictionary with keys
        [min_level, min_level+1, ..., max_level]. The values are tensor with
        shape [height_l, width_l, num_anchors]. The height_l and width_l
        represent the dimension of class logits at l-th level.
      box_targets_dict: ordered dictionary with keys
        [min_level, min_level+1, ..., max_level]. The values are tensor with
        shape [height_l, width_l, num_anchors * 4]. The height_l and
        width_l represent the dimension of bounding box regression output at
        l-th level.
      num_positives: scalar tensor storing number of positives in an image.
    """
    gt_box_list = box_list.BoxList(gt_boxes)
    anchor_box_list = box_list.BoxList(self._anchors.boxes)

    # cls_weights, box_weights are not used
    cls_targets, _, box_targets, _, matches = self._target_assigner.assign(
        anchor_box_list, gt_box_list, gt_labels)

    # class labels start from 1 and the background class = -1
    cls_targets -= 1
    cls_targets = tf.cast(cls_targets, tf.int32)

    # Unpack labels.
    cls_targets_dict = self._unpack_labels(cls_targets)
    box_targets_dict = self._unpack_labels(box_targets)
    num_positives = tf.reduce_sum(
        tf.cast(tf.not_equal(matches.match_results, -1), tf.float32))

    return cls_targets_dict, box_targets_dict, num_positives
//  [python]
def _get_coarse_value(input_img):
  """Maps input with values from `{1, 0, -1}` to a coarse value in `{[0, 1], -1}`.

  Args:
    input_img: Tensor with values from `{1, 0, -1}`, where `1` is positive, `0`
      is negative, and `-1` is uncertain.

  Returns:
    A scalar tensor with value in the range `[0, 1]` indicating the ratio of
    positives to the total `positives / (positives + negatives)`, or `-1` if
    `(positives + negatives)` is `0`.
  """
  positives = tf.math.count_nonzero(tf.equal(input_img, 1))
  negatives = tf.math.count_nonzero(tf.equal(input_img, 0))
  total_pos_neg = tf.math.add(positives, negatives)
  return tf.cond(
      tf.equal(total_pos_neg, 0), lambda: tf.constant(-1.0, dtype=tf.float32),
      lambda: tf.to_float(tf.math.divide(positives, total_pos_neg)))
//  [python]
def _get_coarse_value(input_img):
  """Maps input with values from `{1, 0, -1}` to a coarse value in `{[0, 1], -1}`.

  Args:
    input_img: Tensor with values from `{1, 0, -1}`, where `1` is positive, `0`
      is negative, and `-1` is uncertain.

  Returns:
    A scalar tensor with value in the range `[0, 1]` indicating the ratio of
    positives to the total `positives / (positives + negatives)`, or `-1` if
    `(positives + negatives)` is `0`.
  """
  positives = tf.math.count_nonzero(tf.equal(input_img, 1))
  negatives = tf.math.count_nonzero(tf.equal(input_img, 0))
  total_pos_neg = tf.math.add(positives, negatives)
  return tf.cond(
      tf.equal(total_pos_neg, 0), lambda: tf.constant(-1.0, dtype=tf.float32),
      lambda: tf.cast(tf.math.divide(positives, total_pos_neg)))
//  [python]
def multilayer_perceptron(n_examples, input_shape, output_scaler=1.):
  """Builds a single hidden layer Bayesian feedforward network.

  Args:
    n_examples: Number of examples in training set.
    input_shape: tf.TensorShape.
    output_scaler: Float to scale mean predictions. Training is faster and more
      stable when both the inputs and outputs are normalized. To not affect
      metrics such as RMSE and NLL, the outputs need to be scaled back
      (de-normalized, but the mean doesn't matter), using output_scaler.

  Returns:
    tf.keras.Model.
  """
  p_fn, q_fn = mean_field_fn(empirical_bayes=True)
  def normalized_kl_fn(q, p, _):
    return q.kl_divergence(p) / tf.cast(n_examples, tf.float32)

  inputs = tf.keras.layers.Input(shape=input_shape)
  hidden = tfp.layers.DenseLocalReparameterization(
      50,
      activation='relu',
      kernel_prior_fn=p_fn,
      kernel_posterior_fn=q_fn,
      bias_prior_fn=p_fn,
      bias_posterior_fn=q_fn,
      kernel_divergence_fn=normalized_kl_fn,
      bias_divergence_fn=normalized_kl_fn)(inputs)
  loc = tfp.layers.DenseLocalReparameterization(
      1,
      activation=None,
      kernel_prior_fn=p_fn,
      kernel_posterior_fn=q_fn,
      bias_prior_fn=p_fn,
      bias_posterior_fn=q_fn,
      kernel_divergence_fn=normalized_kl_fn,
      bias_divergence_fn=normalized_kl_fn)(hidden)
  loc = tf.keras.layers.Lambda(lambda x: x * output_scaler)(loc)
  scale = tfp.layers.VariableLayer(
      shape=(), initializer=tf.keras.initializers.Constant(-3.))(loc)
  scale = tf.keras.layers.Activation('softplus')(scale)
  outputs = tf.keras.layers.Lambda(lambda x: ed.Normal(loc=x[0], scale=x[1]))(
      (loc, scale))
  return tf.keras.Model(inputs=inputs, outputs=outputs)
//  [python]
def __init__(self, shape, name=None):
        """Takes input in uint8 format which is cast to float32 and divided by 255
        before passing it to the model.

        On GPU this ensures lower data transfer times.

        Parameters
        ----------
        shape: [int]
            shape of the tensor.
        name: str
            name of the underlying placeholder
        """

        super().__init__(tf.placeholder(
            tf.uint8, [None] + list(shape), name=name))
        self._shape = shape
        self._output = tf.cast(super().get(), tf.float32) / 255.0
//  [python]
def _convert_parameter(param, dtype):
  try:
    return tf.convert_to_tensor(param, dtype=dtype)
  except ValueError:
    try:
      return t.cast(param, dtype)
    except ValueError:
      return t.cast(param(), dtype)
//  [python]
def _convert_parameter(param, dtype):
  try:
    return tf.convert_to_tensor(param, dtype=dtype)
  except ValueError:
    try:
      return tf.cast(param, dtype)
    except ValueError:
      return tf.cast(param(), dtype)
//  [python]
def warmup_lr(self, init_lr, global_step, warmup_step, num_gpu):
        def warmup(end_lr, global_step, warmup_step):
            start_lr = end_lr * 0.1
            global_step = t.cast(global_step, t.float32)
            return start_lr + (end_lr - start_lr) * global_step / warmup_step

        def decay(start_lr, global_step, num_gpu):
            lr = tf.train.piecewise_constant(global_step,
                                             boundaries=[np.int64(self.cfgs.DECAY_STEP[0] // num_gpu),
                                                         np.int64(self.cfgs.DECAY_STEP[1] // num_gpu),
                                                         np.int64(self.cfgs.DECAY_STEP[2] // num_gpu)],
                                             values=[start_lr, start_lr / 10., start_lr / 100., start_lr / 1000.])
            return lr

        return tf.cond(tf.less_equal(global_step, warmup_step),
                       true_fn=lambda: warmup(init_lr, global_step, warmup_step),
                       false_fn=lambda: decay(init_lr, global_step, num_gpu))
//  [python]
def warmup_lr(self, init_lr, global_step, warmup_step, num_gpu):
        def warmup(end_lr, global_step, warmup_step):
            start_lr = end_lr * 0.1
            global_step = tf.cast(global_step, tf.float32)
            return start_lr + (end_lr - start_lr) * global_step / warmup_step

        def decay(start_lr, global_step, num_gpu):
            lr = tf.train.piecewise_constant(global_step,
                                             boundaries=[np.int64(self.cfgs.DECAY_STEP[0] // num_gpu),
                                                         np.int64(self.cfgs.DECAY_STEP[1] // num_gpu),
                                                         np.int64(self.cfgs.DECAY_STEP[2] // num_gpu)],
                                             values=[start_lr, start_lr / 10., start_lr / 100., start_lr / 1000.])
            return lr

        return tf.cond(tf.less_equal(global_step, warmup_step),
                       true_fn=lambda: warmup(init_lr, global_step, warmup_step),
                       false_fn=lambda: decay(init_lr, global_step, num_gpu))
//  [python]
def warmup_and_cosine_lr(self, init_lr, global_step, warmup_step, decay_steps, alpha=1e-6):
        def warmup(end_lr, global_step, warmup_step):
            start_lr = end_lr * 0.1
            global_step = t.cast(global_step, t.float32)
            return start_lr + (end_lr - start_lr) * global_step / warmup_step

        def cosine_lr(init_lr, global_step, decay_steps, alpha=0.0):
            return tf.train.cosine_decay(learning_rate=init_lr,
                                         global_step=global_step - warmup_step,
                                         decay_steps=decay_steps - warmup_step,
                                         alpha=alpha)

        return tf.cond(tf.less_equal(global_step, warmup_step),
                       true_fn=lambda: warmup(init_lr, global_step, warmup_step),
                       false_fn=lambda: cosine_lr(init_lr, global_step, decay_steps, alpha))
//  [python]
def warmup_and_cosine_lr(self, init_lr, global_step, warmup_step, decay_steps, alpha=1e-6):
        def warmup(end_lr, global_step, warmup_step):
            start_lr = end_lr * 0.1
            global_step = tf.cast(global_step, tf.float32)
            return start_lr + (end_lr - start_lr) * global_step / warmup_step

        def cosine_lr(init_lr, global_step, decay_steps, alpha=0.0):
            return tf.train.cosine_decay(learning_rate=init_lr,
                                         global_step=global_step - warmup_step,
                                         decay_steps=decay_steps - warmup_step,
                                         alpha=alpha)

        return tf.cond(tf.less_equal(global_step, warmup_step),
                       true_fn=lambda: warmup(init_lr, global_step, warmup_step),
                       false_fn=lambda: cosine_lr(init_lr, global_step, decay_steps, alpha))
//  [python]
def main():
    """Create the model and start the evaluation process."""
    args = get_arguments()
    
    # Prepare image.
    img = tf.image.decode_jpeg(tf.read_file(args.img_path), channels=3)
    # Convert RGB to BGR.
    img_r, img_g, img_b = tf.split(axis=2, num_or_size_splits=3, value=img)
    img = tf.concat(axis=2, values=[img_b, img_g, img_r])
    img = tf.to_float(img)
    # Extract mean.
    img -= IMG_MEAN 
    
    # Create network.
    net = DeepLabResNetModel({'data': tf.expand_dims(img, dim=0)}, is_training=False, num_classes=args.num_classes)

    # Which variables to load.
    restore_var = tf.global_variables()

    # Predictions.
    raw_output = net.layers['fc1_voc12']
    raw_output_up = tf.image.resize_bilinear(raw_output, tf.shape(img)[0:2,])
    raw_output_up = tf.argmax(raw_output_up, dimension=3)
    pred = tf.expand_dims(raw_output_up, dim=3)

    
    # Set up TF session and initialize variables. 
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    sess = tf.Session(config=config)
    init = tf.global_variables_initializer()
    
    sess.run(init)
    
    # Load weights.
    loader = tf.train.Saver(var_list=restore_var)
    load(loader, sess, args.model_weights)
    
    # Perform inference.
    preds = sess.run(pred)
    
    msk = decode_labels(preds, num_classes=args.num_classes)
    im = Image.fromarray(msk[0])
    if not os.path.exists(args.save_dir):
        os.makedirs(args.save_dir)
    im.save(args.save_dir + 'mask.png')
    
    print('The output file has been saved to {}'.format(args.save_dir + 'mask.png'))
//  [python]
def main():
    """Create the model and start the evaluation process."""
    args = get_arguments()
    
    # Prepare image.
    img = tf.image.decode_jpeg(tf.read_file(args.img_path), channels=3)
    # Convert RGB to BGR.
    img_r, img_g, img_b = tf.split(axis=2, num_or_size_splits=3, value=img)
    img = tf.cast(tf.concat(axis=2, values=[img_b, img_g, img_r]), dtype=tf.float32)
    # Extract mean.
    img -= IMG_MEAN 
    
    # Create network.
    net = DeepLabResNetModel({'data': tf.expand_dims(img, dim=0)}, is_training=False, num_classes=args.num_classes)

    # Which variables to load.
    restore_var = tf.global_variables()

    # Predictions.
    raw_output = net.layers['fc1_voc12']
    raw_output_up = tf.image.resize_bilinear(raw_output, tf.shape(img)[0:2,])
    raw_output_up = tf.argmax(raw_output_up, dimension=3)
    pred = tf.expand_dims(raw_output_up, dim=3)

    
    # Set up TF session and initialize variables. 
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    sess = tf.Session(config=config)
    init = tf.global_variables_initializer()
    
    sess.run(init)
    
    # Load weights.
    loader = tf.train.Saver(var_list=restore_var)
    load(loader, sess, args.model_weights)
    
    # Perform inference.
    preds = sess.run(pred)
    
    msk = decode_labels(preds, num_classes=args.num_classes)
    im = Image.fromarray(msk[0])
    if not os.path.exists(args.save_dir):
        os.makedirs(args.save_dir)
    im.save(args.save_dir + 'mask.png')
    
    print('The output file has been saved to {}'.format(args.save_dir + 'mask.png'))
//  [python]
async def predict(
    image: fastapi.UploadFile = fastapi.File(default=None),
) -> fastapi.responses.JSONResponse:
  """Predicts objects in the uploaded image.

  Args:
    image: Image from which to generate predictions.

  Returns:
    A JSON encoded list of detections.
  """
  image_data = await image.read()
  try:
    p_image = PIL.Image.open(io.BytesIO(image_data))
  except (OSError, PIL.UnidentifiedImageError):
    return fastapi.responses.JSONResponse(
        content={'message': 'Could not open image_data as an image.'},
        status_code=400,
    )  # Bad Request

  try:
    tf_image = tf.image.resize(
        p_image, (HEIGHT, WIDTH), method=tf.image.ResizeMethod.AREA
    )
    image_cp = tf.cast(tf_image, tf.uint8)
    image = app_utils.preprocess_image(image_cp)
    detections = list(
        map(
            lambda model: app_utils.perform_detection(model, image),
            model_manager.detection_fns,
        )
    )

    json_dump = json.dumps(
        {'predictions': detections}, cls=app_utils.NumpyEncoder
    )
    return fastapi.responses.JSONResponse(content=json_dump)

  except TypeError:
    return fastapi.responses.JSONResponse(
        content={'message': 'Image data is not in the correct format.'},
        status_code=422,
    )
//  [python]
def main(args):
    g = tf.Graph()      # A new graph
    with g.as_default():
        with tf.Session() as sess:
            # Building graph.
            image_data = tf.placeholder(tf.int32, name='input_image')
            height = tf.placeholder(tf.int32, name='height')
            width = tf.placeholder(tf.int32, name='width')

            # Reshape data
            image = tf.reshape(image_data, [height, width, 3])

            processed_image = utils.mean_image_subtraction(
                image, [123.68, 116.779, 103.939])                    # Preprocessing image
            batched_image = tf.expand_dims(processed_image, 0)        # Add batch dimension
            generated_image = model.net(batched_image, training=False)
            casted_image = tf.to_int32(generated_image)
            # Remove batch dimension
            squeezed_image = tf.squeeze(casted_image, [0])
            cropped_image = tf.slice(squeezed_image, [0, 0, 0], [height, width, 3])
            # stylized_image = tf.image.encode_jpeg(squeezed_image, name='output_image')
            stylized_image_data = tf.reshape(cropped_image, [-1], name='output_image')

            # Restore model variables.
            saver = tf.train.Saver(tf.global_variables(), write_version=tf.train.SaverDef.V1)
            sess.run([tf.global_variables_initializer(), tf.local_variables_initializer()])
            # Use absolute path.
            model_file = os.path.abspath(args.model_file)
            saver.restore(sess, model_file)

            if args.is_debug:
                content_file = '/Users/Lex/Desktop/t.jpg'
                generated_file = '/Users/Lex/Desktop/xwz-stylized.jpg'

                with open(generated_file, 'wb') as img:
                    image_bytes = tf.read_file(content_file)
                    input_array, decoded_image = sess.run([
                        tf.reshape(tf.image.decode_jpeg(image_bytes, channels=3), [-1]),
                        tf.image.decode_jpeg(image_bytes, channels=3)])

                    start_time = time.time()
                    img.write(sess.run(tf.image.encode_jpeg(tf.to_int32(cropped_image), tf.uint8), feed_dict={
                              image_data: input_array,
                              height: decoded_image.shape[0],
                              width: decoded_image.shape[1]}))
                    end_time = time.time()

                    tf.logging.info('Elapsed time: %fs' % (end_time - start_time))
            else:
                output_graph_def = tf.graph_util.convert_variables_to_constants(
                    sess, sess.graph_def, output_node_names=['output_image'])

                with tf.gfile.FastGFile('/Users/Lex/Desktop/' + args.model_name + '.pb', mode='wb') as f:
                    f.write(output_graph_def.SerializeToString())
//  [python]
def main(args):
    g = tf.Graph()      # A new graph
    with g.as_default():
        with tf.Session() as sess:
            # Building graph.
            image_data = tf.placeholder(tf.int32, name='input_image')
            height = tf.placeholder(tf.int32, name='height')
            width = tf.placeholder(tf.int32, name='width')

            # Reshape data
            image = tf.reshape(image_data, [height, width, 3])

            processed_image = utils.mean_image_subtraction(
                image, [123.68, 116.779, 103.939])                    # Preprocessing image
            batched_image = tf.expand_dims(processed_image, 0)        # Add batch dimension
            generated_image = model.net(batched_image, training=False)
            casted_image = tf.cast(generated_image, tf.int32)
            # Remove batch dimension
            squeezed_image = tf.squeeze(casted_image, [0])
            cropped_image = tf.slice(squeezed_image, [0, 0, 0], [height, width, 3])
            # stylized_image = tf.image.encode_jpeg(squeezed_image, name='output_image')
            stylized_image_data = tf.reshape(cropped_image, [-1], name='output_image')

            # Restore model variables.
            saver = tf.train.Saver(tf.global_variables(), write_version=tf.train.SaverDef.V1)
            sess.run([tf.global_variables_initializer(), tf.local_variables_initializer()])
            # Use absolute path.
            model_file = os.path.abspath(args.model_file)
            saver.restore(sess, model_file)

            if args.is_debug:
                content_file = '/Users/Lex/Desktop/t.jpg'
                generated_file = '/Users/Lex/Desktop/xwz-stylized.jpg'

                with open(generated_file, 'wb') as img:
                    image_bytes = tf.read_file(content_file)
                    input_array, decoded_image = sess.run([
                        tf.reshape(tf.image.decode_jpeg(image_bytes, channels=3), [-1]),
                        tf.image.decode_jpeg(image_bytes, channels=3)])

                    start_time = time.time()
                    img.write(sess.run(tf.image.encode_jpeg(tf.cast(cropped_image, tf.uint8)), feed_dict={
                              image_data: input_array,
                              height: decoded_image.shape[0],
                              width: decoded_image.shape[1]}))
                    end_time = time.time()

                    tf.logging.info('Elapsed time: %fs' % (end_time - start_time))
            else:
                output_graph_def = tf.graph_util.convert_variables_to_constants(
                    sess, sess.graph_def, output_node_names=['output_image'])

                with tf.gfile.FastGFile('/Users/Lex/Desktop/' + args.model_name + '.pb', mode='wb') as f:
                    f.write(output_graph_def.SerializeToString())
//  [python]
def loss(logits, labels):
    labels = tf.to_int32(labels)
    cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(
        logits=logits, labels=labels, name='cross_entropy_per_example')
    cross_entropy_mean = tf.reduce_mean(cross_entropy, name='cross_entropy')
    tf.add_to_collection('losses', cross_entropy_mean)
    losses = tf.get_collection('losses')
    regularization_losses = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)
    total_loss = cross_entropy_mean + LAMBDA * sum(regularization_losses)
    tf.summary.scalar('tl (raw)', total_loss)
    #total_loss = tf.add_n(losses + regularization_losses, name='total_loss')
    loss_averages = tf.train.ExponentialMovingAverage(0.9, name='avg')
    loss_averages_op = loss_averages.apply(losses + [total_loss])
    for l in losses + [total_loss]:
        tf.summary.scalar(l.op.name + ' (raw)', l)
        tf.summary.scalar(l.op.name, loss_averages.average(l))
    with tf.control_dependencies([loss_averages_op]):
        total_loss = tf.identity(total_loss)
    return total_loss
//  [python]
def loss(logits, labels):
    labels = tf.cast(labels, tf.int32)
    cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(
        logits=logits, labels=labels, name='cross_entropy_per_example')
    cross_entropy_mean = tf.reduce_mean(cross_entropy, name='cross_entropy')
    tf.add_to_collection('losses', cross_entropy_mean)
    losses = tf.get_collection('losses')
    regularization_losses = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)
    total_loss = cross_entropy_mean + LAMBDA * sum(regularization_losses)
    tf.summary.scalar('tl (raw)', total_loss)
    #total_loss = tf.add_n(losses + regularization_losses, name='total_loss')
    loss_averages = tf.train.ExponentialMovingAverage(0.9, name='avg')
    loss_averages_op = loss_averages.apply(losses + [total_loss])
    for l in losses + [total_loss]:
        tf.summary.scalar(l.op.name + ' (raw)', l)
        tf.summary.scalar(l.op.name, loss_averages.average(l))
    with tf.control_dependencies([loss_averages_op]):
        total_loss = tf.identity(total_loss)
    return total_loss
//  [python]
def get_train_op(FLAGS, total_loss, grads_and_vars=None):
  global_step = tf.train.get_or_create_global_step()

  # increase the learning rate linearly
  if FLAGS.warmup_steps > 0:
    warmup_lr = (tf.to_float(global_step)
                 / tf.to_float(FLAGS.warmup_steps)
                 * FLAGS.learning_rate)
  else:
    warmup_lr = 0.0

  # decay the learning rate
  if FLAGS.decay_method == "poly":
    decay_lr = tf.train.polynomial_decay(
        FLAGS.learning_rate,
        global_step=global_step - FLAGS.warmup_steps,
        decay_steps=FLAGS.train_steps - FLAGS.warmup_steps,
        end_learning_rate=FLAGS.learning_rate * FLAGS.min_lr_ratio)
  elif FLAGS.decay_method == "cos":
    decay_lr = tf.train.cosine_decay(
        FLAGS.learning_rate,
        global_step=global_step - FLAGS.warmup_steps,
        decay_steps=FLAGS.train_steps - FLAGS.warmup_steps,
        alpha=FLAGS.min_lr_ratio)
  else:
    raise ValueError(FLAGS.decay_method)

  learning_rate = tf.where(global_step < FLAGS.warmup_steps,
                           warmup_lr, decay_lr)

  if (FLAGS.weight_decay > 0 and not FLAGS.use_tpu and
      FLAGS.num_core_per_host > 1):
    raise ValueError("Do not support `weight_decay > 0` with multi-gpu "
                     "training so far.")

  if FLAGS.weight_decay == 0:
    optimizer = tf.train.AdamOptimizer(
        learning_rate=learning_rate,
        epsilon=FLAGS.adam_epsilon)
  else:
    optimizer = AdamWeightDecayOptimizer(
        learning_rate=learning_rate,
        epsilon=FLAGS.adam_epsilon,
        exclude_from_weight_decay=["LayerNorm", "layer_norm", "bias"],
        weight_decay_rate=FLAGS.weight_decay)

  if FLAGS.use_tpu:
    optimizer = tf.contrib.tpu.CrossShardOptimizer(optimizer)

  if grads_and_vars is None:
    grads_and_vars = optimizer.compute_gradients(total_loss)
  gradients, variables = zip(*grads_and_vars)
  clipped, gnorm = tf.clip_by_global_norm(gradients, FLAGS.clip)

  if getattr(FLAGS, "lr_layer_decay_rate", 1.0) != 1.0:
    n_layer = 0
    for i in range(len(clipped)):
      m = re.search(r"model/transformer/layer_(\d+?)/", variables[i].name)
      if not m: continue
      n_layer = max(n_layer, int(m.group(1)) + 1)

    for i in range(len(clipped)):
      for l in range(n_layer):
        if "model/transformer/layer_{}/".format(l) in variables[i].name:
          abs_rate = FLAGS.lr_layer_decay_rate ** (n_layer - 1 - l)
          clipped[i] *= abs_rate
          tf.logging.info("Apply mult {:.4f} to layer-{} grad of {}".format(
              abs_rate, l, variables[i].name))
          break

  train_op = optimizer.apply_gradients(
      zip(clipped, variables), global_step=global_step)

  # Manually increment `global_step` for AdamWeightDecayOptimizer
  if FLAGS.weight_decay > 0:
    new_global_step = global_step + 1
    train_op = tf.group(train_op, [global_step.assign(new_global_step)])

  return train_op, learning_rate, gnorm
//  [python]
def get_train_op(FLAGS, total_loss, grads_and_vars=None):
  global_step = tf.train.get_or_create_global_step()

  # increase the learning rate linearly
  if FLAGS.warmup_steps > 0:
    warmup_lr = (tf.cast(global_step, tf.float32)
                 / tf.cast(FLAGS.warmup_steps, tf.float32)
                 * FLAGS.learning_rate)
  else:
    warmup_lr = 0.0

  # decay the learning rate
  if FLAGS.decay_method == "poly":
    decay_lr = tf.train.polynomial_decay(
        FLAGS.learning_rate,
        global_step=global_step - FLAGS.warmup_steps,
        decay_steps=FLAGS.train_steps - FLAGS.warmup_steps,
        end_learning_rate=FLAGS.learning_rate * FLAGS.min_lr_ratio)
  elif FLAGS.decay_method == "cos":
    decay_lr = tf.train.cosine_decay(
        FLAGS.learning_rate,
        global_step=global_step - FLAGS.warmup_steps,
        decay_steps=FLAGS.train_steps - FLAGS.warmup_steps,
        alpha=FLAGS.min_lr_ratio)
  else:
    raise ValueError(FLAGS.decay_method)

  learning_rate = tf.where(global_step < FLAGS.warmup_steps,
                           warmup_lr, decay_lr)

  if (FLAGS.weight_decay > 0 and not FLAGS.use_tpu and
      FLAGS.num_core_per_host > 1):
    raise ValueError("Do not support `weight_decay > 0` with multi-gpu "
                     "training so far.")

  if FLAGS.weight_decay == 0:
    optimizer = tf.train.AdamOptimizer(
        learning_rate=learning_rate,
        epsilon=FLAGS.adam_epsilon)
  else:
    optimizer = AdamWeightDecayOptimizer(
        learning_rate=learning_rate,
        epsilon=FLAGS.adam_epsilon,
        exclude_from_weight_decay=["LayerNorm", "layer_norm", "bias"],
        weight_decay_rate=FLAGS.weight_decay)

  if FLAGS.use_tpu:
    optimizer = tf.contrib.tpu.CrossShardOptimizer(optimizer)

  if grads_and_vars is None:
    grads_and_vars = optimizer.compute_gradients(total_loss)
  gradients, variables = zip(*grads_and_vars)
  clipped, gnorm = tf.clip_by_global_norm(gradients, FLAGS.clip)

  if getattr(FLAGS, "lr_layer_decay_rate", 1.0) != 1.0:
    n_layer = 0
    for i in range(len(clipped)):
      m = re.search(r"model/transformer/layer_(\d+?)/", variables[i].name)
      if not m: continue
      n_layer = max(n_layer, int(m.group(1)) + 1)

    for i in range(len(clipped)):
      for l in range(n_layer):
        if "model/transformer/layer_{}/".format(l) in variables[i].name:
          abs_rate = FLAGS.lr_layer_decay_rate ** (n_layer - 1 - l)
          clipped[i] *= abs_rate
          tf.logging.info("Apply mult {:.4f} to layer-{} grad of {}".format(
              abs_rate, l, variables[i].name))
          break

  train_op = optimizer.apply_gradients(
      zip(clipped, variables), global_step=global_step)

  # Manually increment `global_step` for AdamWeightDecayOptimizer
  if FLAGS.weight_decay > 0:
    new_global_step = global_step + 1
    train_op = tf.group(train_op, [global_step.assign(new_global_step)])

  return train_op, learning_rate, gnorm
//  [python]
def serve(
      self, images: tf.Tensor) -> Mapping[str, tf.Tensor]:
    """Casts an image tensor to float and runs inference.

    Args:
      images: A uint8 tf.Tensor of shape [batch_size, None, None, None,
        num_channels].

    Returns:
      A dictionary holding segmentation outputs.
    """
    with tf.device('cpu:0'):
      images = tf.to_float(images)

    outputs = self.inference_step(images)
    output_key = 'logits' if self.params.task.model.head.output_logits else 'probs'

    return {output_key: outputs['logits']}
//  [python]
def serve(
      self, images: tf.Tensor) -> Mapping[str, tf.Tensor]:
    """Casts an image tensor to float and runs inference.

    Args:
      images: A uint8 tf.Tensor of shape [batch_size, None, None, None,
        num_channels].

    Returns:
      A dictionary holding segmentation outputs.
    """
    with tf.device('cpu:0'):
      images = tf.cast(images, dtype=tf.float32)

    outputs = self.inference_step(images)
    output_key = 'logits' if self.params.task.model.head.output_logits else 'probs'

    return {output_key: outputs['logits']}
//  [python]
def decode(serialized_example):
  """Parses an image and label from the given `serialized_example`."""
  features = tf.parse_single_example(
      serialized_example,
      # Defaults are not specified since both keys are required.
      features={
          'image_raw': tf.FixedLenFeature([], tf.string),
          'label': tf.FixedLenFeature([], tf.int64),
      })

  # Convert from a scalar string tensor (whose single string has
  # length mnist.IMAGE_PIXELS) to a uint8 tensor with shape
  # [mnist.IMAGE_PIXELS].
  image = tf.decode_raw(features['image_raw'], tf.uint8)
  image.set_shape((mnist.IMAGE_PIXELS))

  # Convert label from a scalar uint8 tensor to an int32 scalar.
  label = tf.to_int32(features['label'])

  return image, label
//  [python]
def decode(serialized_example):
  """Parses an image and label from the given `serialized_example`."""
  features = tf.parse_single_example(
      serialized_example,
      # Defaults are not specified since both keys are required.
      features={
          'image_raw': tf.FixedLenFeature([], tf.string),
          'label': tf.FixedLenFeature([], tf.int64),
      })

  # Convert from a scalar string tensor (whose single string has
  # length mnist.IMAGE_PIXELS) to a uint8 tensor with shape
  # [mnist.IMAGE_PIXELS].
  image = tf.decode_raw(features['image_raw'], tf.uint8)
  image.set_shape((mnist.IMAGE_PIXELS))

  # Convert label from a scalar uint8 tensor to an int32 scalar.
  label = tf.cast(features['label'], tf.int32)

  return image, label
//  [python]
def normalize(image, label):
  """Convert `image` from [0, 255] -> [-0.5, 0.5] floats."""
  image = t.cast(image, tf.float32) * (1. / 255) - 0.5
  return image, label
//  [python]
def normalize(image, label):
  """Convert `image` from [0, 255] -> [-0.5, 0.5] floats."""
  image = tf.cast(image, tf.float32) * (1. / 255) - 0.5
  return image, label
//  [python]
def hf_compute_loss(self, logits, labels):
        # upsample logits to the images' original size
        # `labels` is of shape (batch_size, height, width)
        label_interp_shape = shape_list(labels)[1:]

        upsampled_logits = tf.image.resize(logits, size=label_interp_shape, method="bilinear")
        # compute weighted loss
        loss_fct = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction="none")

        def masked_loss(real, pred):
            unmasked_loss = loss_fct(real, pred)
            mask = tf.to_float(real != self.config.semantic_loss_ignore_index)
            masked_loss = unmasked_loss * mask
            # Reduction strategy in the similar spirit with
            # https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_tf_utils.py#L210
            reduced_masked_loss = tf.reduce_sum(masked_loss) / tf.reduce_sum(mask)
            return tf.reshape(reduced_masked_loss, (1,))

        return masked_loss(labels, upsampled_logits)
//  [python]
def hf_compute_loss(self, logits, labels):
        # upsample logits to the images' original size
        # `labels` is of shape (batch_size, height, width)
        label_interp_shape = shape_list(labels)[1:]

        upsampled_logits = tf.image.resize(logits, size=label_interp_shape, method="bilinear")
        # compute weighted loss
        loss_fct = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction="none")

        def masked_loss(real, pred):
            unmasked_loss = loss_fct(real, pred)
            mask = tf.cast(real != self.config.semantic_loss_ignore_index, dtype=unmasked_loss.dtype)
            masked_loss = unmasked_loss * mask
            # Reduction strategy in the similar spirit with
            # https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_tf_utils.py#L210
            reduced_masked_loss = tf.reduce_sum(masked_loss) / tf.reduce_sum(mask)
            return tf.reshape(reduced_masked_loss, (1,))

        return masked_loss(labels, upsampled_logits)
//  [python]
def call(self, encoder_hidden_states, training: bool = False):
        batch_size = shape_list(encoder_hidden_states[-1])[0]

        all_hidden_states = ()
        for encoder_hidden_state, mlp in zip(encoder_hidden_states, self.mlps):
            if self.config.reshape_last_stage is False and len(shape_list(encoder_hidden_state)) == 3:
                height = tf.math.sqrt(tf.cast(shape_list(encoder_hidden_state)[1], tf.float32))
                height = width = tf.cast(height, tf.int32)
                encoder_hidden_state = tf.reshape(encoder_hidden_state, (batch_size, height, width, -1))

            # unify channel dimension
            encoder_hidden_state = tf.transpose(encoder_hidden_state, perm=[0, 2, 3, 1])
            height = shape_list(encoder_hidden_state)[1]
            width = shape_list(encoder_hidden_state)[2]
            encoder_hidden_state = mlp(encoder_hidden_state)
            encoder_hidden_state = tf.reshape(encoder_hidden_state, (batch_size, height, width, -1))

            # upsample
            temp_state = tf.transpose(encoder_hidden_states[0], perm=[0, 2, 3, 1])
            upsample_resolution = shape_list(temp_state)[1:-1]
            encoder_hidden_state = tf.image.resize(encoder_hidden_state, size=upsample_resolution, method="bilinear")
            all_hidden_states += (encoder_hidden_state,)

        hidden_states = self.linear_fuse(tf.concat(all_hidden_states[::-1], axis=-1))
        hidden_states = self.batch_norm(hidden_states, training=training)
        hidden_states = self.activation(hidden_states)
        hidden_states = self.dropout(hidden_states, training=training)

        # logits of shape (batch_size, height/4, width/4, num_labels)
        logits = self.classifier(hidden_states)

        return logits
//  [python]
def call(
        self,
        hidden_states: tf.Tensor,
        height: int,
        width: int,
        output_attentions: bool = False,
        training: bool = False,
    ) -> Union[tf.Tensor, Tuple[tf.Tensor, tf.Tensor]]:
        batch_size = shape_list(hidden_states)[0]
        num_channels = shape_list(hidden_states)[2]

        query_layer = self.transpose_for_scores(self.query(hidden_states))

        if self.sr_ratio > 1:
            # Reshape to (batch_size, height, width, num_channels)
            hidden_states = tf.reshape(hidden_states, (batch_size, height, width, num_channels))
            # Apply sequence reduction
            hidden_states = self.sr(hidden_states)
            # Reshape back to (batch_size, seq_len, num_channels)
            hidden_states = tf.reshape(hidden_states, (batch_size, -1, num_channels))
            hidden_states = self.layer_norm(hidden_states)

        key_layer = self.transpose_for_scores(self.key(hidden_states))
        value_layer = self.transpose_for_scores(self.value(hidden_states))

        # Take the dot product between "query" and "key" to get the raw attention scores.
        attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)

        scale = tf.to_float(self.sqrt_att_head_size)
        attention_scores = tf.divide(attention_scores, scale)

        # Normalize the attention scores to probabilities.
        attention_probs = stable_softmax(logits=attention_scores, axis=-1)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(attention_probs, training=training)

        context_layer = tf.matmul(attention_probs, value_layer)

        context_layer = tf.transpose(context_layer, perm=[0, 2, 1, 3])
        # (batch_size, seq_len_q, all_head_size)
        context_layer = tf.reshape(context_layer, (batch_size, -1, self.all_head_size))

        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
        return outputs
//  [python]
def call(
        self,
        hidden_states: tf.Tensor,
        height: int,
        width: int,
        output_attentions: bool = False,
        training: bool = False,
    ) -> Union[tf.Tensor, Tuple[tf.Tensor, tf.Tensor]]:
        batch_size = shape_list(hidden_states)[0]
        num_channels = shape_list(hidden_states)[2]

        query_layer = self.transpose_for_scores(self.query(hidden_states))

        if self.sr_ratio > 1:
            # Reshape to (batch_size, height, width, num_channels)
            hidden_states = tf.reshape(hidden_states, (batch_size, height, width, num_channels))
            # Apply sequence reduction
            hidden_states = self.sr(hidden_states)
            # Reshape back to (batch_size, seq_len, num_channels)
            hidden_states = tf.reshape(hidden_states, (batch_size, -1, num_channels))
            hidden_states = self.layer_norm(hidden_states)

        key_layer = self.transpose_for_scores(self.key(hidden_states))
        value_layer = self.transpose_for_scores(self.value(hidden_states))

        # Take the dot product between "query" and "key" to get the raw attention scores.
        attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)

        scale = tf.cast(self.sqrt_att_head_size, dtype=attention_scores.dtype)
        attention_scores = tf.divide(attention_scores, scale)

        # Normalize the attention scores to probabilities.
        attention_probs = stable_softmax(logits=attention_scores, axis=-1)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(attention_probs, training=training)

        context_layer = tf.matmul(attention_probs, value_layer)

        context_layer = tf.transpose(context_layer, perm=[0, 2, 1, 3])
        # (batch_size, seq_len_q, all_head_size)
        context_layer = tf.reshape(context_layer, (batch_size, -1, self.all_head_size))

        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
        return outputs
//  [python]
def __init__(self, x, y, num_classes, dtype=tf.float32, learn_rate=1e-3):
    x_shape = x.get_shape()
    x_size = 1
    for ss in x_shape[1:]:
      x_size *= int(ss)
    x = tf.reshape(x, [-1, x_size])
    w_class = weight_variable(
        [x_size, num_classes],
        init_method='truncated_normal',
        dtype=tf.float32,
        init_param={'stddev': 0.01},
        name='w_class')
    b_class = weight_variable(
        [num_classes],
        init_method='constant',
        init_param={'val': 0.0},
        name='b_class')
    logits = tf.matmul(x, w_class) + b_class
    xent = tf.nn.sparse_softmax_cross_entropy_with_logits(
        logits=logits, labels=y)
    xent = tf.reduce_mean(xent, name='xent')
    cost = xent
    cost += self._decay()
    self._cost = cost
    self._inputs = x
    self._labels = y
    self._train_op = tf.train.AdamOptimizer(learn_rate).minimize(
        cost, var_list=[w_class, b_class])
    correct = tf.equal(tf.argmax(logits, axis=1), y)
    self._acc = tf.reduce_mean(tf.to_float(correct))
    self._prediction = tf.nn.softmax(logits)
//  [python]
def __init__(self, x, y, num_classes, dtype=tf.float32, learn_rate=1e-3):
    x_shape = x.get_shape()
    x_size = 1
    for ss in x_shape[1:]:
      x_size *= int(ss)
    x = tf.reshape(x, [-1, x_size])
    w_class = weight_variable(
        [x_size, num_classes],
        init_method='truncated_normal',
        dtype=tf.float32,
        init_param={'stddev': 0.01},
        name='w_class')
    b_class = weight_variable(
        [num_classes],
        init_method='constant',
        init_param={'val': 0.0},
        name='b_class')
    logits = tf.matmul(x, w_class) + b_class
    xent = tf.nn.sparse_softmax_cross_entropy_with_logits(
        logits=logits, labels=y)
    xent = tf.reduce_mean(xent, name='xent')
    cost = xent
    cost += self._decay()
    self._cost = cost
    self._inputs = x
    self._labels = y
    self._train_op = tf.train.AdamOptimizer(learn_rate).minimize(
        cost, var_list=[w_class, b_class])
    correct = tf.equal(tf.argmax(logits, axis=1), y)
    self._acc = tf.reduce_mean(tf.cast(correct, dtype))
    self._prediction = tf.nn.softmax(logits)